| """Learned three-action baseline with explicit uncertainty abstention.""" |
| from dataclasses import asdict |
| import json |
| from pathlib import Path |
| import re |
| import torch |
| from .actions import Action, Decision, Kind |
| from .features import encode, KINDS |
| from .model import PointerPolicy |
|
|
|
|
| class LearnedPolicy: |
| def __init__(self, checkpoint, quantized=False, confidence_threshold=.85): |
| root = Path(checkpoint) |
| self.model = PointerPolicy.from_pretrained(root).eval() |
| self.vocab = json.loads((root/'vocab.json').read_text(encoding='utf-8')) |
| self.temperatures = json.loads((root/'calibration.json').read_text())['temperatures'] |
| self.threshold = confidence_threshold |
| if quantized: |
| self.model = torch.ao.quantization.quantize_dynamic(self.model,{torch.nn.Linear},dtype=torch.qint8) |
|
|
| @torch.inference_mode() |
| def predict(self, goal, state, ticket): |
| row = dict(goal=goal,elements=[asdict(e) for e in state.elements]) |
| inputs,_,_,maps = encode([row],self.vocab) |
| if not maps[0]: |
| return Decision(ticket,Action(Kind.ASK_USER,('No eligible DOM target; a fallback is required.',))) |
| a,t = self.model(*inputs) |
| ap = (a/self.temperatures[0]).softmax(-1)[0] |
| tp = (t/self.temperatures[1]).softmax(-1)[0] |
| ai,ti = int(ap.argmax()),int(tp.argmax()) |
| ac,tc = float(ap[ai]),float(tp[ti]) |
| if min(ac,tc) < self.threshold: |
| return Decision(ticket,Action(Kind.ASK_USER,('The policy is uncertain about this action or target.',)),ac,tc) |
| kind = Kind(KINDS[ai]) |
| element = state.elements[maps[0][ti]] |
| args = (element.ref,) |
| if kind in {Kind.TYPE,Kind.SELECT}: |
| |
| literals = re.findall(r'"([^"\n]*)"',goal) |
| if len(literals) != 1: |
| return Decision(ticket,Action(Kind.ASK_USER,('This baseline needs one quoted value to copy.',)),ac,tc) |
| args += (literals[0],) |
| return Decision(ticket,Action(kind,args),ac,tc) |
|
|
