File size: 2,116 Bytes
795f737 | 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 33 34 35 36 37 38 39 40 41 42 43 44 45 | """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}:
# Literal copying, not inferred website/workflow logic. General span prediction is pending.
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)
|