"""CPU decision inference for the 17,601-parameter standalone policy."""
import argparse,json
from pathlib import Path
import torch
from safetensors.torch import load_file
from huggingface_hub import hf_hub_download
from structured_policy import StructuredPolicy,encode
ap=argparse.ArgumentParser();ap.add_argument('--request',default='example.json')
ap.add_argument('--weights',default=None);a=ap.parse_args()
path=a.weights or hf_hub_download('guilindev/pacman-decision-tiny','model.safetensors')
model=StructuredPolicy().eval();model.load_state_dict(load_file(path))
r=json.load(open(a.request,encoding='utf-8'));keys=list(r['questions']['move']['criteria'])
assert 2<=len(keys)<=4
x=torch.tensor([encode(r['state'],keys)]);valid=torch.ones(1,len(keys),dtype=torch.bool)
with torch.inference_mode():p=model(x,valid).softmax(-1)[0].tolist()
print(json.dumps({'choice':keys[max(range(len(p)),key=p.__getitem__)],'probabilities':dict(zip(keys,p))},indent=2))