| """Decision making: menu-mediated, model-agnostic. |
| |
| The design lesson from testing tiny models on real screenshots: a 256M VLM can |
| *read* a button label from pixels but cannot reliably reason, and a 0.5B text model |
| writes perfect JSON but picks the wrong widget when handed noisy free-form text. |
| Both do much better when the harness hands them a **short numbered menu** and asks |
| for a number. That is what this module builds, and it keeps the model's output space |
| finite no matter how small the model is. |
| |
| Brains shipped here: |
| |
| * :class:`MockBrain` - replays a scripted action list (tests, demos, replays) |
| * :class:`PlannerBrain` - rule-based, no model: clicks the element matching the goal |
| * :class:`ModelBrain` - base for real models: menu, prompt, lenient parsing |
| (see :mod:`tiny_computer_agent.models` for the text and vision backends) |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| import re |
| from dataclasses import dataclass |
| from typing import Any, Iterator, Mapping, Protocol, Sequence |
|
|
| from tiny_computer_agent.actions import ACTION_KINDS, Action, ActionError, parse_action |
| from tiny_computer_agent.decoding import ( |
| brevity_ladder, |
| choice_token_ids, |
| digit_token_ids, |
| majority_choice, |
| pair_token_ids, |
| read_choice, |
| read_verdict, |
| ) |
| from tiny_computer_agent.observe import Element, Observation, normalise_label |
|
|
|
|
| class BrainError(RuntimeError): |
| """Raised when a brain cannot produce a valid action.""" |
|
|
|
|
| class GoalUnreachable(BrainError): |
| """Raised by a brain that can decide nothing further, without claiming success. |
| |
| Deliberately distinct from a crash: the harness records the episode as a |
| *failure*, because a brain admitting "I give up" must never be logged as |
| "goal met" - and the harness, not the brain, owns that judgement. |
| """ |
|
|
|
|
| def extract_json_objects(raw: str) -> list[dict[str, Any]]: |
| """Every flat JSON object in a reply, in the order they appear. |
| |
| Regex-based and non-nesting: replies from tiny models are one-level objects. |
| Callers should prefer the **last** object, because a model that echoes its own |
| prompt before answering puts the prompt's example first and its answer last. |
| """ |
| objects: list[dict[str, Any]] = [] |
| for match in re.finditer(r"\{[^{}]*\}", raw or "", re.DOTALL): |
| try: |
| payload = json.loads(match.group(0)) |
| except json.JSONDecodeError: |
| continue |
| if isinstance(payload, dict): |
| objects.append(payload) |
| return objects |
|
|
|
|
| def extract_json_object(raw: str) -> dict[str, Any]: |
| """The decision out of a reply: the *last* JSON object, not the first.""" |
| objects = extract_json_objects(raw) |
| if not objects: |
| raise BrainError(f"no JSON object in the model reply: {(raw or '')[:200]!r}") |
| return objects[-1] |
|
|
|
|
| @dataclass(frozen=True) |
| class MenuItem: |
| """One numbered choice shown to the model.""" |
|
|
| index: int |
| element: Element |
|
|
| @property |
| def label(self) -> str: |
| return self.element.label |
|
|
|
|
| class DecisionMenu: |
| """A short, numbered, closed choice set built from the element index. |
| |
| Elements are presented in reading order (top to bottom, then left to right), |
| and capped: a 20-item menu is as good as free-form text to a small model. |
| """ |
|
|
| def __init__(self, observation: Observation, *, max_items: int = 12) -> None: |
| ordered = sorted(observation.elements, key=lambda e: (e.y, e.x))[:max_items] |
| self.items: tuple[MenuItem, ...] = tuple( |
| MenuItem(index=i + 1, element=e) for i, e in enumerate(ordered) |
| ) |
|
|
| def __len__(self) -> int: |
| return len(self.items) |
|
|
| def __bool__(self) -> bool: |
| return bool(self.items) |
|
|
| def __iter__(self) -> Iterator[MenuItem]: |
| return iter(self.items) |
|
|
| def item(self, value: Any) -> MenuItem | None: |
| """Look up a 1-based index, tolerating the strings models actually send.""" |
| try: |
| number = int(float(str(value).strip())) |
| except (TypeError, ValueError): |
| return None |
| if 1 <= number <= len(self.items): |
| return self.items[number - 1] |
| return None |
|
|
| def find(self, label: str) -> Element | None: |
| """Resolve a label against the menu (normalised, like Observation.find).""" |
| wanted = normalise_label(label) |
| if not wanted: |
| return None |
| for item in self.items: |
| if normalise_label(item.label) == wanted: |
| return item.element |
| for item in self.items: |
| if wanted in normalise_label(item.label): |
| return item.element |
| return None |
|
|
| def mentioned_in(self, text: str) -> tuple[MenuItem, ...]: |
| """Menu entries whose label appears in the text, in order of first mention.""" |
| haystack = normalise_label(text) |
| found: list[tuple[int, MenuItem]] = [] |
| for item in self.items: |
| needle = normalise_label(item.label) |
| if needle and (position := haystack.find(needle)) != -1: |
| found.append((position, item)) |
| return tuple(item for _, item in sorted(found, key=lambda pair: pair[0])) |
|
|
| def render(self) -> str: |
| return "\n".join( |
| f" {item.index}. {item.label} (at {item.element.center[0]},{item.element.center[1]})" |
| for item in self.items |
| ) |
|
|
|
|
| def _recover_from_text(payload: Mapping[str, Any], menu: DecisionMenu) -> Action | None: |
| """Salvage a decision from a payload that named the right thing in the wrong field. |
| |
| Observed for real: a 0.5B planner replied ``{"kind": "click", "index": 0, |
| "reason": "Save"}`` - the correct answer, in the wrong place. If the index is |
| unusable but exactly one menu label is mentioned anywhere in the payload, take |
| the label. It is still only ever a label the detector saw in this frame, so the |
| deny list, the dry run and the step budget all continue to apply. |
| """ |
| text = json.dumps(payload, ensure_ascii=False) |
| mentioned = menu.mentioned_in(text) |
| if len({item.label for item in mentioned}) == 1: |
| item = mentioned[0] |
| return Action( |
| "click", label=item.label, |
| reason=f"recovered {item.label!r} from the model's reply: {text[:80]}", |
| ) |
| return None |
|
|
|
|
| def _action_from_payload(payload: Mapping[str, Any], menu: DecisionMenu) -> Action | None: |
| """Interpret one JSON object, tolerating the ways small models mangle the shape.""" |
| kind = payload.get("kind") |
| if isinstance(kind, str) and kind.strip().lower() not in ACTION_KINDS: |
| kind = None |
|
|
| if payload.get("index") is not None: |
| item = menu.item(payload["index"]) |
| if item is not None: |
| return Action("click", label=item.label, reason=f"menu index {payload['index']}") |
| recovered = _recover_from_text(payload, menu) |
| if recovered is not None: |
| return recovered |
| raise BrainError( |
| f"the model chose index {payload['index']!r}, which is not in the " |
| f"menu of {len(menu)} items" |
| ) |
|
|
| if kind in (None, "click") and "label" in payload: |
| element = menu.find(str(payload["label"])) |
| if element is None: |
| raise BrainError(f"the model named {payload['label']!r}, which is not on screen") |
| return Action("click", label=element.label, reason=f"the model named {element.label!r}") |
|
|
| if kind is not None: |
| try: |
| return parse_action(dict(payload, kind=kind)) |
| except ActionError as exc: |
| |
| return _recover_from_text(payload, menu) or (_raise_brain_error(exc)) |
|
|
| return _recover_from_text(payload, menu) |
|
|
|
|
| def _raise_brain_error(exc: Exception) -> Action: |
| raise BrainError(f"the model proposed an invalid action: {exc}") from exc |
|
|
|
|
| _DONE = re.compile(r"\b(done|finished|complete[d]?)\b", re.IGNORECASE) |
|
|
|
|
| def parse_decision(raw: str, observation: Observation, *, max_menu_items: int = 12) -> Action: |
| """Turn a small model's reply into a validated Action - leniently, but safely. |
| |
| Tried in order: |
| |
| 1. the **last** JSON object in the reply (``index``, ``label`` or a full action) |
| 2. a label the reply names (most specific evidence of intent) |
| 3. a bare menu number |
| 4. a "done" word |
| |
| A reply that mentions three or more menu entries is rejected as an echo of the |
| prompt rather than a decision - the failure mode of the weakest models, and the |
| one most likely to make an agent click something arbitrary. |
| """ |
| text = " ".join((raw or "").split()) |
| if not text: |
| raise BrainError("the model returned an empty reply") |
| menu = DecisionMenu(observation, max_items=max_menu_items) |
|
|
| for payload in reversed(extract_json_objects(text)): |
| action = _action_from_payload(payload, menu) |
| if action is not None: |
| return action |
|
|
| mentioned = menu.mentioned_in(text) |
| if len(mentioned) >= 3: |
| raise BrainError( |
| f"the reply mentions {len(mentioned)} menu entries, so it looks like an echo " |
| f"of the prompt rather than a decision: {text[:120]!r}" |
| ) |
| if mentioned: |
| item = mentioned[0] |
| note = f"the reply names {item.label!r}" |
| if len(mentioned) > 1: |
| note += f" (ambiguous: {len(mentioned)} labels mentioned; took the first)" |
| return Action("click", label=item.label, reason=note) |
|
|
| number = re.search(r"\b(\d{1,2})\b", text) |
| if number is not None: |
| item = menu.item(number.group(1)) |
| if item is not None: |
| return Action( |
| "click", label=item.label, |
| reason=f"menu index {number.group(1)} from reply {text[:40]!r}", |
| ) |
|
|
| if _DONE.search(text): |
| return Action("done", reason=_DONE.search(text).group(0).lower()) |
|
|
| raise BrainError(f"could not read a decision out of the reply: {text[:120]!r}") |
|
|
|
|
| class Brain(Protocol): |
| """Anything that can turn (goal, observation) into one validated action.""" |
|
|
| name: str |
|
|
| def decide(self, goal: str, observation: Observation) -> Action: ... |
|
|
|
|
| class MockBrain: |
| """Replays a fixed script of actions, exactly once. Deterministic and offline.""" |
|
|
| name = "mock" |
|
|
| def __init__(self, script: Sequence[Action | Mapping[str, Any]]) -> None: |
| self._script: list[Action] = [ |
| a if isinstance(a, Action) else parse_action(a) for a in script |
| ] |
| self._cursor = 0 |
|
|
| @property |
| def remaining(self) -> int: |
| return max(0, len(self._script) - self._cursor) |
|
|
| def decide(self, goal: str, observation: Observation) -> Action: |
| if self._cursor >= len(self._script): |
| return Action("done", reason="script exhausted") |
| action = self._script[self._cursor] |
| self._cursor += 1 |
| return action |
|
|
|
|
| class PlannerBrain: |
| """Rule-based baseline: click the element matching the goal, then stop. |
| |
| Its purpose is threefold - it proves the harness works end to end with no model |
| at all, it is a fallback when a model download is undesirable, and it shows how |
| much of a "computer use" task is really an indexing problem. It waits when the |
| target is not on screen yet, refuses to click the same label twice, and gives up |
| rather than flailing. It will not generalise the way a model can. |
| """ |
|
|
| name = "planner" |
|
|
| def __init__(self, target: str | None = None, *, |
| done_when_text_contains: str | None = None, |
| max_waits: int = 3, wait_seconds: float = 0.5) -> None: |
| self.target = target |
| self.done_when_text_contains = done_when_text_contains |
| self.max_waits = max_waits |
| self.wait_seconds = wait_seconds |
| self._clicked: set[str] = set() |
| self._waits = 0 |
|
|
| def decide(self, goal: str, observation: Observation) -> Action: |
| if self.done_when_text_contains: |
| needle = self.done_when_text_contains.lower() |
| if needle in observation.text.lower(): |
| return Action("done", reason=f"page text contains {self.done_when_text_contains!r}") |
|
|
| target = self.target or self._target_from(goal) |
| if target: |
| element = observation.find(target) |
| if element is None: |
| if self._waits < self.max_waits: |
| self._waits += 1 |
| return Action("wait", seconds=self.wait_seconds, |
| reason=f"{target!r} is not on screen yet") |
| raise GoalUnreachable(f"{target!r} never appeared on screen") |
| key = element.label.strip().lower() |
| if key in self._clicked: |
| return Action("done", reason=f"already clicked {element.label!r}; goal assumed met") |
| self._clicked.add(key) |
| return Action("click", label=element.label, |
| reason=f"matched {target!r} in the element index") |
| raise GoalUnreachable(f"no element name found in the goal {goal!r}") |
|
|
| @staticmethod |
| def _target_from(goal: str) -> str | None: |
| """Best effort: quoted text wins, else the last capitalised word.""" |
| quoted = re.search(r"[\"\']([^\"\']{2,40})[\"\']", goal) |
| if quoted: |
| return quoted.group(1) |
| words = [w for w in goal.split() if w[:1].isupper()] |
| return words[-1].strip(".,!?") if words else None |
|
|
|
|
| class ModelBrain: |
| """Shared plumbing for model-backed brains: menu, prompt, lenient parsing. |
| |
| Subclasses implement :meth:`_generate`. ``requires_index`` marks brains that |
| cannot decide without an element index (a text-only planner): for those, an |
| empty screen is a hard stop rather than a prompt. |
| """ |
|
|
| name = "model" |
| requires_index = False |
| |
| |
| supports_constrained = False |
| |
| |
| |
| |
| |
| decision_mode = "json" |
|
|
| def __init__(self, model_id: str, *, max_menu_items: int = 12, |
| max_new_tokens: int = 64, decision_mode: str = "json", |
| retries: int = 3, votes: int = 1) -> None: |
| self.model_id = model_id |
| self.max_menu_items = max_menu_items |
| self.max_new_tokens = max_new_tokens |
| self.decision_mode = decision_mode |
| self.retries = max(1, retries) |
| self.votes = max(1, votes) |
| self.last_choice: int | None = None |
| self.echoes_refused = 0 |
| self.last_prompt: str | None = None |
| self.last_reply: str | None = None |
| |
| |
| self.history: list[str] = [] |
| self._last_fingerprint: tuple | None = None |
|
|
| @staticmethod |
| def _fingerprint(observation: Observation) -> tuple: |
| return (tuple(element.label for element in observation.elements), |
| observation.text[:200]) |
|
|
| def build_prompt(self, goal: str, menu: DecisionMenu, observation: Observation, |
| *, history: Sequence[str] = (), screen_changed: bool = True, |
| brevity: int = 0) -> str: |
| """A short prompt with a strict output shape. |
| |
| Deliberately contains placeholders (``<number>``) rather than worked |
| examples: the weakest models echo the prompt back, and an echoed |
| placeholder fails validation instead of being mistaken for a decision. |
| |
| ``brevity`` is the retry ladder: echoing is driven by how much text the |
| model has to copy, so each step upwards drops a section - first the screen |
| text and the history, then everything but the menu. |
| """ |
| if self.decision_mode == "choice": |
| return self._choice_prompt(goal, menu, observation, |
| history=history, brevity=brevity) |
| lines = [ |
| "You are controlling a computer. Choose exactly one action.", |
| "", |
| "Reply with one JSON object and nothing else, in this shape:", |
| ' {"kind": "click", "index": <number from the list>, "reason": "<brief>"}', |
| ' {"kind": "done", "reason": "<brief>"}', |
| "Use done only when the goal is already met.", |
| "", |
| f"Goal: {goal}", |
| ] |
| if observation.text: |
| lines += ["", f"Text on screen: {observation.text[:300]}"] |
| lines += ["", f"Clickable items ({len(menu)}):"] |
| lines.append(menu.render() if menu else " (none detected)") |
| if history: |
| lines += ["", "Already attempted - do NOT repeat these, they did not work:"] |
| lines += [f" - {entry}" for entry in history] |
| if not screen_changed: |
| lines += ["", "The screen did not change after your last action. " |
| "Choose a DIFFERENT action."] |
| return "\n".join(lines) |
|
|
| def _choice_prompt(self, goal: str, menu: DecisionMenu, observation: Observation, |
| *, history: Sequence[str] = (), brevity: int = 0) -> str: |
| """The constrained-mode prompt: one number, and as little text as possible. |
| |
| There is no JSON shape to copy and no worked example - the whole point of |
| the masked path is that the model has nothing to echo. |
| """ |
| lines = [ |
| "You are controlling a computer. Reply with ONE NUMBER and nothing else.", |
| "0 means the goal is already met. Any other number is the item to click.", |
| "", |
| f"Goal: {goal}", |
| ] |
| if brevity == 0 and observation.text: |
| lines += ["", f"Text on screen: {observation.text[:200]}"] |
| lines += ["", f"Items ({len(menu)}):", menu.render() if menu else " (none detected)"] |
| if brevity == 0 and history: |
| lines += ["", "Already tried, and it did not work:"] |
| lines += [f" - {entry}" for entry in history] |
| return "\n".join(lines) |
|
|
| def decide(self, goal: str, observation: Observation) -> Action: |
| menu = DecisionMenu(observation, max_items=self.max_menu_items) |
| if self.requires_index and not menu: |
| raise GoalUnreachable( |
| "no clickable elements on screen - is screen capture and OCR working?" |
| ) |
| fingerprint = self._fingerprint(observation) |
| screen_changed = (self._last_fingerprint is None |
| or fingerprint != self._last_fingerprint) |
|
|
| |
| |
| |
| if self.history: |
| previous = self.history[-1].split(" -> ")[0] |
| self.history[-1] = ( |
| f"{previous} -> " |
| + ("the screen changed" if screen_changed else "the screen did NOT change") |
| ) |
|
|
| self.last_prompt = self.build_prompt( |
| goal, menu, observation, history=tuple(self.history), |
| screen_changed=screen_changed, |
| ) |
| if self.decision_mode == "rank" and self.supports_constrained: |
| action = self._decide_by_ranking(goal, menu, observation) |
| elif self.decision_mode == "scan" and self.supports_constrained: |
| action = self._decide_by_scanning(goal, menu, observation) |
| elif self.decision_mode == "choice" and self.supports_constrained: |
| action = self._decide_by_choice(goal, menu, observation) |
| else: |
| self.last_reply = self._generate(self.last_prompt, observation) |
| action = parse_decision(self.last_reply, observation, |
| max_menu_items=self.max_menu_items) |
| self.history.append(action.describe()) |
| del self.history[:-3] |
| self._last_fingerprint = fingerprint |
| return action |
|
|
| def _decide_by_choice(self, goal: str, menu: DecisionMenu, |
| observation: Observation) -> Action: |
| """Ask for a menu number with every other token masked out. |
| |
| The echo failure cannot happen here: the model is unable to emit the prompt |
| back, only one of the numbers we allow. When even that fails, the prompt |
| gets shorter (echoing is driven by prompt length) and, if ``votes`` is set, |
| the answer is taken by majority over several samples. |
| """ |
| if not menu: |
| raise GoalUnreachable("no clickable items to choose from") |
| tokenizer = self._tokenizer_for_choice() |
| allowed = choice_token_ids(tokenizer, len(menu)) |
| if not allowed: |
| allowed = digit_token_ids(tokenizer) |
| if not allowed: |
| raise BrainError( |
| "this tokenizer cannot express a menu number as a single token, so " |
| "the choice cannot be constrained; use decision_mode='json'" |
| ) |
|
|
| last_error: BrainError | None = None |
| for brevity in brevity_ladder(self.retries): |
| self.last_prompt = self.build_prompt( |
| goal, menu, observation, history=tuple(self.history), brevity=brevity) |
| replies = [ |
| self._generate_constrained(self.last_prompt, observation, allowed, 4, |
| sample=index > 0) |
| for index in range(self.votes) |
| ] |
| self.last_reply = replies[0] if replies else "" |
| choices = [read_choice(reply, len(menu)) for reply in replies] |
| number = majority_choice(choices) |
| if number is None: |
| self.echoes_refused += sum(1 for reply in replies |
| if reply and reply.strip() == self.last_prompt.strip()) |
| last_error = BrainError( |
| f"the constrained reply was not a menu number: " |
| f"{self.last_reply[:80]!r} (brevity {brevity}, {len(replies)} vote(s))" |
| ) |
| continue |
| self.last_choice = number |
| if number == 0: |
| return Action("done", reason="the model chose 0: the goal is already met") |
| item = menu.item(number) |
| if item is None: |
| last_error = BrainError(f"the model chose {number}, outside the menu") |
| continue |
| return Action("click", label=item.label, |
| reason=f"constrained choice {number} of {len(menu)}") |
|
|
| raise last_error or BrainError("the model produced no constrained choice") |
|
|
| def _scan_prompt(self, goal: str, label: str) -> str: |
| """The prompt for one binary question: does *this* item do the job? |
| |
| Short on purpose. Every word a tiny model is given is a word it can copy |
| instead of answering, and there is nothing here worth copying back. |
| """ |
| return "\n".join([ |
| "You are controlling a computer. Answer with one word: yes or no.", |
| "", |
| f"Goal: {goal}", |
| f"Candidate item: {label}", |
| "Does clicking this item achieve the goal?", |
| ]) |
|
|
| def _decide_by_scanning(self, goal: str, menu: DecisionMenu, |
| observation: Observation) -> Action: |
| """Ask one masked yes/no question per menu item, and take the first yes. |
| |
| Why this rather than one N-way choice: a 256M model cannot hold twelve |
| options in mind, but "does this button do what the goal says?" is a |
| question it can answer. Each answer is one token out of two, so the mask |
| makes the reply space as small as it can possibly be - and the scan stops |
| at the first yes, so it is usually cheaper than it looks. |
| """ |
| if not menu: |
| raise GoalUnreachable("no clickable items to choose from") |
| tokenizer = self._tokenizer_for_choice() |
| yes_ids, no_ids = pair_token_ids(tokenizer) |
| if not yes_ids or not no_ids: |
| raise BrainError( |
| "this tokenizer cannot express yes/no as single tokens, so the scan " |
| "cannot be constrained; use decision_mode='json'" |
| ) |
| allowed = sorted(set(yes_ids) | set(no_ids)) |
| scanned: list[tuple[str, str | None]] = [] |
| for item in menu: |
| self.last_prompt = self._scan_prompt(goal, item.label) |
| self.last_reply = self._generate_constrained( |
| self.last_prompt, observation, allowed, 1, sample=False) |
| verdict = read_verdict(self.last_reply) |
| scanned.append((item.label, verdict)) |
| if verdict is None: |
| |
| |
| raise BrainError( |
| f"the masked reply to a yes/no question was {self.last_reply[:40]!r}" |
| ) |
| if verdict == "yes": |
| return Action("click", label=item.label, |
| reason=f"the model said yes to {item.label!r} " |
| f"({len(scanned)} of {len(menu)} scanned)") |
| self.last_reply = "" |
| raise BrainError( |
| f"the model said no to all {len(menu)} item(s): " |
| f"{[label for label, _ in scanned][:6]}" |
| ) |
|
|
| def _decide_by_ranking(self, goal: str, menu: DecisionMenu, |
| observation: Observation) -> Action: |
| """Ask every item, then click the one the model liked most. |
| |
| ``scan`` stops at the first yes, which quietly rewards position. Here the |
| model is asked about every candidate and the winner is the largest |
| yes-minus-no log-odds margin - the model's own opinion, read off the |
| distribution instead of off a sampled token. Costs N short generations, |
| and it is the cheapest way to find out whether the weights can do it at all. |
| """ |
| if not menu: |
| raise GoalUnreachable("no clickable items to choose from") |
| tokenizer = self._tokenizer_for_choice() |
| yes_ids, no_ids = pair_token_ids(tokenizer) |
| if not yes_ids or not no_ids: |
| raise BrainError( |
| "this tokenizer cannot express yes/no as single tokens, so the ranking " |
| "cannot be constrained; use decision_mode='json'" |
| ) |
| scored: list[tuple[Any, float]] = [] |
| for item in menu: |
| self.last_prompt = self._scan_prompt(goal, item.label) |
| margin = self._score_pair(self.last_prompt, observation, yes_ids, no_ids) |
| scored.append((item, margin)) |
| scored.sort(key=lambda pair: pair[1], reverse=True) |
| best, margin = scored[0] |
| self.last_reply = f"margin {margin:+.2f}" |
| if margin <= 0: |
| raise BrainError( |
| f"the model preferred no item: best was {best.label!r} at {margin:+.2f}; " |
| f"runner-up {scored[1][0].label!r} at {scored[1][1]:+.2f}" |
| if len(scored) > 1 else |
| f"the model preferred no item (only {best.label!r}, {margin:+.2f})" |
| ) |
| return Action("click", label=best.label, |
| reason=f"highest yes/no margin {margin:+.2f} of {len(scored)} item(s)") |
|
|
| def _score_pair(self, prompt: str, observation: Observation, |
| yes_ids: Sequence[int], no_ids: Sequence[int]) -> float: |
| """How much more likely "yes" is than "no". Backends implement this.""" |
| raise NotImplementedError("a constrained backend implements _score_pair") |
|
|
| def _tokenizer_for_choice(self) -> Any: |
| """The tokenizer used to build the mask. Subclasses load it lazily.""" |
| raise NotImplementedError("a constrained backend provides its tokenizer") |
|
|
| def _generate_constrained(self, prompt: str, observation: Observation, |
| allowed_ids: Sequence[int], max_new_tokens: int, |
| *, sample: bool = False) -> str: |
| raise NotImplementedError("a constrained backend implements this") |
|
|
| @property |
| def loaded(self) -> bool: |
| """True once the weights are in memory - they load lazily, on the first decide().""" |
| return getattr(self, "_model", None) is not None |
|
|
| def _generate(self, prompt: str, observation: Observation) -> str: |
| raise NotImplementedError( |
| "subclass ModelBrain and implement _generate(prompt, observation)" |
| ) |
|
|
