Optimize CPU feature encoding and add bounded batch inference with measured parity

README.md

@@ -11,6 +11,8 @@ library_name: baim
 
 # Devils Agent / BAIM — experimental research checkpoints
 
+Runtime v0.0.2: **2.38x faster median synthetic CPU prediction**, bounded microbatch inference, and unchanged checkpoint weights. See [measured results and limits](docs/RUNTIME_V002.md).
+
 **Research prototype, not a production-ready general browser agent.** Four custom
 checkpoints are stored under `models/`. They require the accompanying Python code;
 this repository is not a standard Transformers `AutoModel` or hosted-inference

baim/features.py

@@ -1,62 +1,84 @@
-"""Small vocabulary and deterministic candidate retrieval, without site rules."""
-from collections import Counter
-import re
-import torch
-
-KINDS = ('C', 'T', 'O')
-ROLES = {'button', 'link', 'textbox', 'searchbox', 'combobox', 'checkbox', 'radio', 'listbox', 'spinbutton'}
-
-
-def words(text):
-    return re.findall(r'\w+', text.casefold(), flags=re.UNICODE)
-
-
-def fit_vocab(rows, limit=4096):
-    counts = Counter()
-    for row in rows:
-        counts.update(words(row['goal']))
-        for element in row['elements']:
-            counts.update(words(element['role'] + ' ' + element['name']))
-    return {'<pad>':0, '<unk>':1, **{token:i+2 for i, (token, _) in enumerate(counts.most_common(limit-2))}}
-
-
-def lexical(goal, element):
-    goal_words, name_words = set(words(goal)), set(words(element['name']))
-    overlap = len(goal_words & name_words)
-    return [overlap / max(1,len(name_words)), overlap / max(1,len(goal_words)),
-            float(element['name'].casefold() in goal.casefold()),
-            float(element.get('visible', True)), float(element.get('enabled', True))]
-
-
-def candidates(goal, elements, limit=40):
-    eligible = [i for i, e in enumerate(elements) if e.get('visible', True) and e.get('enabled', True)
-                and not e.get('sensitive', False) and e['role'] in ROLES]
-    return sorted(eligible, key=lambda i: (-lexical(goal,elements[i])[0], i))[:limit]
-
-
-def encode(rows, vocab, max_tokens=24, max_candidates=40):
-    # Training padding uses batch size; inference uses only the surviving candidates.
-    maps = [candidates(row['goal'],row['elements'],max_candidates) for row in rows]
-    count = max(1, max(map(len,maps),default=0))
-    goal = torch.zeros((len(rows),max_tokens),dtype=torch.long)
-    element = torch.zeros((len(rows),count,max_tokens),dtype=torch.long)
-    features = torch.zeros((len(rows),count,5))
-    mask = torch.zeros((len(rows),count),dtype=torch.bool)
-    targets = torch.full((len(rows),),-100,dtype=torch.long)
-    actions = torch.full((len(rows),),-100,dtype=torch.long)
-    def tokens(text):
-        return [vocab.get(token,1) for token in words(text)[:max_tokens]]
-    for batch,row in enumerate(rows):
-        ids = tokens(row['goal'])
-        goal[batch,:len(ids)] = torch.tensor(ids,dtype=torch.long)
-        if row.get('action') in KINDS:
-            actions[batch] = KINDS.index(row['action'])
-        for j,index in enumerate(maps[batch]):
-            e = row['elements'][index]
-            ids = tokens(e['role'] + ' ' + e['name'])
-            element[batch,j,:len(ids)] = torch.tensor(ids,dtype=torch.long)
-            features[batch,j] = torch.tensor(lexical(row['goal'],e))
-            mask[batch,j] = True
-            if row.get('target') == index:
-                targets[batch] = j
-    return (goal,element,features,mask),actions,targets,maps
+"""Small vocabulary and deterministic candidate retrieval, without site rules."""
+from collections import Counter
+import re
+import torch
+
+KINDS = ('C', 'T', 'O')
+ROLES = {'button', 'link', 'textbox', 'searchbox', 'combobox', 'checkbox', 'radio', 'listbox', 'spinbutton'}
+
+
+def words(text):
+    return re.findall(r'\w+', text.casefold(), flags=re.UNICODE)
+
+
+def fit_vocab(rows, limit=4096):
+    counts = Counter()
+    for row in rows:
+        counts.update(words(row['goal']))
+        for element in row['elements']:
+            counts.update(words(element['role'] + ' ' + element['name']))
+    return {'<pad>':0, '<unk>':1, **{token:i+2 for i, (token, _) in enumerate(counts.most_common(limit-2))}}
+
+
+def lexical(goal, element):
+    goal_words, name_words = set(words(goal)), set(words(element['name']))
+    overlap = len(goal_words & name_words)
+    return [overlap / max(1,len(name_words)), overlap / max(1,len(goal_words)),
+            float(element['name'].casefold() in goal.casefold()),
+            float(element.get('visible', True)), float(element.get('enabled', True))]
+
+
+def candidates(goal, elements, limit=40):
+    eligible = [i for i, e in enumerate(elements) if e.get('visible', True) and e.get('enabled', True)
+                and not e.get('sensitive', False) and e['role'] in ROLES]
+    return sorted(eligible, key=lambda i: (-lexical(goal,elements[i])[0], i))[:limit]
+
+
+def encode(rows, vocab, max_tokens=24, max_candidates=40):
+    # Tokenize each goal once and construct six tensors per batch, not per element.
+    if max_tokens < 1 or max_candidates < 1:
+        raise ValueError('Token and candidate limits must be positive')
+    maps, prepared = [], []
+    for row in rows:
+        goal_words = words(row['goal'])
+        goal_set, folded = set(goal_words), row['goal'].casefold()
+        eligible = []
+        for i, e in enumerate(row['elements']):
+            if not (e.get('visible', True) and e.get('enabled', True)) or e.get('sensitive', False) or e['role'] not in ROLES:
+                continue
+            name_set = set(words(e['name']))
+            overlap = len(goal_set & name_set)
+            feature = [overlap/max(1,len(name_set)), overlap/max(1,len(goal_set)),
+                       float(e['name'].casefold() in folded), float(e.get('visible',True)), float(e.get('enabled',True))]
+            eligible.append((i, feature))
+        eligible.sort(key=lambda item: (-item[1][0], item[0]))
+        selected = eligible[:max_candidates]
+        maps.append([i for i, _ in selected])
+        prepared.append((goal_words, selected))
+    count = max(1, max(map(len, maps), default=0))
+    goals, elements, features, masks, actions, targets = [], [], [], [], [], []
+    def tokens(tokens):
+        ids = [vocab.get(token, 1) for token in tokens[:max_tokens]]
+        return ids + [0]*(max_tokens-len(ids))
+    for row, (goal_words, selected) in zip(rows, prepared):
+        goals.append(tokens(goal_words))
+        ids, feats = [], []
+        target = -100
+        for j, (i, feature) in enumerate(selected):
+            e = row['elements'][i]
+            ids.append(tokens(words(e['role']+' '+e['name'])))
+            feats.append(feature)
+            if row.get('target') == i:
+                target = j
+        padding = count-len(selected)
+        elements.append(ids + [[0]*max_tokens for _ in range(padding)])
+        features.append(feats + [[0.0]*5 for _ in range(padding)])
+        masks.append([True]*len(selected) + [False]*padding)
+        actions.append(KINDS.index(row['action']) if row.get('action') in KINDS else -100)
+        targets.append(target)
+    batch = len(rows)
+    inputs = (torch.tensor(goals,dtype=torch.long).reshape(batch,max_tokens),
+              torch.tensor(elements,dtype=torch.long).reshape(batch,count,max_tokens),
+              torch.tensor(features,dtype=torch.float32).reshape(batch,count,5),
+              torch.tensor(masks,dtype=torch.bool).reshape(batch,count))
+    return inputs, torch.tensor(actions,dtype=torch.long), torch.tensor(targets,dtype=torch.long), maps

baim/policy.py

@@ -1,5 +1,4 @@
 """Learned three-action baseline with explicit uncertainty abstention."""
-from dataclasses import asdict
 import json
 from pathlib import Path
 import re
@@ -21,19 +20,34 @@ class LearnedPolicy:
 
     @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.',)))
+        return self.predict_batch([(goal, state, ticket)])[0]
+
+    @torch.inference_mode()
+    def predict_batch(self, requests):
+        """Bounded CPU microbatch. Every decision retains its own authority ticket."""
+        if len(requests) > 64:
+            raise ValueError('A microbatch may contain at most 64 requests')
+        if not requests:
+            return []
+        rows = [dict(goal=goal,elements=[dict(role=e.role,name=e.name,visible=e.visible,
+                    enabled=e.enabled,sensitive=e.sensitive) for e in state.elements])
+                for goal,state,_ in requests]
+        inputs,_,_,maps = encode(rows,self.vocab)
         a,t = self.model(*inputs)
-        ap = (a/self.temperatures[0]).softmax(-1)[0]
-        tp = (t/self.temperatures[1]).softmax(-1)[0]
+        ap = (a/self.temperatures[0]).softmax(-1)
+        tp = (t/self.temperatures[1]).softmax(-1)
+        return [self._decision(goal,state,ticket,mapping,ap[i],tp[i])
+                for i,((goal,state,ticket),mapping) in enumerate(zip(requests,maps))]
+
+    def _decision(self, goal, state, ticket, mapping, ap, tp):
+        if not mapping:
+            return Decision(ticket,Action(Kind.ASK_USER,('No eligible DOM target; a fallback is required.',)))
         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]]
+        element = state.elements[mapping[ti]]
         args = (element.ref,)
         if kind in {Kind.TYPE,Kind.SELECT}:
             # Literal copying, not inferred website/workflow logic. General span prediction is pending.

docs/RUNTIME_V002.md

@@ -0,0 +1,28 @@
+# Runtime v0.0.2
+
+This release optimizes inference code; checkpoint weights and model intelligence are unchanged.
+
+Feature encoding tokenizes each goal once, reuses lexical scores for retrieval, and
+constructs tensors once per batch instead of repeatedly assigning tiny tensors.
+Policy conversion no longer recursively copies unused DOM fields.
+
+`LearnedPolicy.predict_batch([(goal, state, ticket), ...])` supports at most 64
+independent requests. Each returned decision retains its own ticket and observed
+target references. Confidence abstention, sensitive-target exclusion, and value-copy
+restrictions remain in effect. Empty batches return an empty list.
+
+Measured on the development Windows CPU with two Torch threads and the mean
+checkpoint: median single-request prediction 2.103 ms before / 0.884 ms after
+(2.38x); p95 3.304 / 1.327 ms. A 16-request batch took 6.253 ms median, about
+0.391 ms per request. These are synthetic CPU policy timings, not browser-task
+latency or target Linux VPS results. Loading, network, and browser time are excluded.
+
+All encoded inputs, targets, and candidate maps were bit-identical across 1,440
+validation/test/novel-wording examples. Individual action/ticket parity was checked
+on 160 requests; batch actions matched individual actions on those requests.
+Novel-wording raw step match remains 55.83%; this release does not fix generalization.
+
+Evidence: `reports/runtime-optimization.json`. Run `python -m unittest discover -s tests`
+for runtime regression tests. The existing `baim.bench_policy` command can measure
+the current feature encoder independently. Compare against Hub revision
+`795f73703a1e33b82636b2167d7d9507581efd4d` for the prior implementation.

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "baim"
-version = "0.0.1"
+version = "0.0.2"
 description = "Experimental CPU-first browser action intelligence"
 requires-python = ">=3.12"
 dependencies = []

reports/runtime-optimization.json

@@ -0,0 +1,34 @@
+{
+  "scope": "Windows CPU, two Torch threads, synthetic data; unchanged model weights",
+  "splits": {
+    "validation": {
+      "rows": 480,
+      "encoding_bit_exact": true,
+      "raw_step_match": 1.0
+    },
+    "test": {
+      "rows": 480,
+      "encoding_bit_exact": true,
+      "raw_step_match": 1.0
+    },
+    "novel_wording": {
+      "rows": 480,
+      "encoding_bit_exact": true,
+      "raw_step_match": 0.5583333333333333
+    }
+  },
+  "decision_parity_rows": 160,
+  "old_predict": {
+    "median_ms": 2.1032000076957047,
+    "p95_ms": 3.303700010292232
+  },
+  "new_predict": {
+    "median_ms": 0.8837499772198498,
+    "p95_ms": 1.3269999762997031
+  },
+  "batch16": {
+    "median_ms": 6.252999999560416,
+    "p95_ms": 9.634000016376376
+  },
+  "speedup": 2.37985862733719
+}

tests/test_batch.py

@@ -0,0 +1,43 @@
+import unittest
+from pathlib import Path
+import torch
+from baim.policy import LearnedPolicy
+from baim.features import encode, candidates, lexical
+from baim.state import State, Element
+from baim.actions import Ticket, Kind
+
+class BatchTests(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        torch.set_num_threads(2)
+        cls.policy=LearnedPolicy(Path(__file__).resolve().parents[1]/'models/v000-mean')
+
+    def test_batch_isolates_tickets_refs_and_sensitive_targets(self):
+        requests=[]
+        for i in range(4):
+            elements=() if i==3 else (Element(f'{i}-safe','button','amber birch','0'),
+                Element(f'{i}-secret','textbox','password','1',sensitive=True))
+            state=State(str(i),i,'fixture',elements)
+            ticket=Ticket('session',str(i),0,i,str(i),i,state.state_hash)
+            requests.append(('Click amber birch.',state,ticket))
+        batch=self.policy.predict_batch(requests)
+        for request,decision in zip(requests,batch):
+            self.assertEqual(decision.action,self.policy.predict(*request).action)
+            self.assertEqual(decision.ticket,request[2])
+            if decision.action.kind!=Kind.ASK_USER:
+                self.assertEqual(decision.action.args[0],request[1].elements[0].ref)
+        self.assertEqual(batch[-1].action.kind,Kind.ASK_USER)
+        self.assertEqual(self.policy.predict_batch([]),[])
+        with self.assertRaises(ValueError):self.policy.predict_batch(requests*17)
+
+    def test_encoding_preserves_retrieval_and_lexical_features(self):
+        row={'goal':'Click café', 'elements':[{'role':'button','name':'café'},
+            {'role':'button','name':'café','sensitive':True},
+            {'role':'button','name':'other'}, {'role':'button','name':'café','visible':False}]}
+        inputs,_,_,maps=encode([row],self.policy.vocab)
+        self.assertEqual(maps[0],candidates(row['goal'],row['elements']))
+        for i,index in enumerate(maps[0]):
+            self.assertTrue(torch.equal(inputs[2][0,i],torch.tensor(lexical(row['goal'],row['elements'][index]))))
+        self.assertEqual(encode([],self.policy.vocab)[0][0].shape[0],0)
+
+if __name__=='__main__':unittest.main()