File size: 979 Bytes
1501ad1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 | """Run one native typed-choice request from the downloaded model directory."""
import argparse,json
from pathlib import Path
ap=argparse.ArgumentParser();ap.add_argument('--backend',choices=['torch','mlx'],default='torch');ap.add_argument('--device')
ap.add_argument('--request',type=Path);ap.add_argument('--output',type=Path);a=ap.parse_args()
root=Path(__file__).resolve().parent
request=json.loads((a.request or root/'example_request.json').read_text(encoding='utf-8'))
if a.backend=='mlx':
from laya_mlx.agent import Agent
agent=Agent(str(root),dtype='float16',device=a.device or 'gpu',compile=False,cache_prompts=False)
else:
import torch
from laya.agent import Agent
torch.set_num_threads(4);agent=Agent(str(root),device=a.device or 'cuda',compile=False)
result=agent.predict(request['state'],request['questions']);serialized=json.dumps(result,ensure_ascii=False,indent=2)
if a.output:a.output.write_text(serialized,encoding='utf-8')
print(serialized)
|