| """Real (tiny) model backends, small enough to run on a laptop CPU. |
| |
| Two of them, because they trade off differently: |
| |
| * :class:`TextPlannerBrain` - a ~0.5B *text* model over the OCR index. No vision, |
| no big download, and the element index does all the grounding, which is why it |
| works at this size at all. |
| * :class:`VLMBrain` - a tiny *vision* model that sees the screenshot itself |
| (256M-class models can read a button label from pixels, but reason poorly, so the |
| menu still carries most of the decision). |
| |
| Both inherit menu building, prompt shape and lenient parsing from |
| :class:`~tiny_computer_agent.brain.ModelBrain`, so swapping weights never changes the |
| harness contract. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import os |
| from typing import Any |
|
|
| from tiny_computer_agent.brain import BrainError, ModelBrain |
| from tiny_computer_agent.decoding import MenuMask |
| from tiny_computer_agent.observe import Observation |
|
|
| DEFAULT_TEXT_MODEL = "Qwen/Qwen2.5-0.5B-Instruct" |
| DEFAULT_VISION_MODEL = "HuggingFaceTB/SmolVLM-256M-Instruct" |
|
|
|
|
| def _pair_margin(logits: Any, yes_ids: Any, no_ids: Any) -> float: |
| """log P(yes) - log P(no) for one position, over the whole vocabulary.""" |
| import torch |
|
|
| yes = torch.logsumexp(logits[list(yes_ids)], dim=0) |
| no = torch.logsumexp(logits[list(no_ids)], dim=0) |
| return float((yes - no).item()) |
|
|
|
|
| def _masked_generation(model: Any, allowed_ids: Any, eos_token_id: Any) -> Any: |
| """A transformers ``LogitsProcessorList`` holding the menu mask. |
| |
| Imported here, not at module scope, so the harness still imports with no |
| transformers installed at all. |
| """ |
| try: |
| from transformers import LogitsProcessorList |
| except ImportError: |
| raise BrainError( |
| "constrained decoding needs transformers: pip install 'tiny-computer-agent[vision]'" |
| ) |
| return LogitsProcessorList([MenuMask(allowed_ids, eos_token_id=eos_token_id)]) |
|
|
|
|
| def resolve_device(prefer: str | None = None) -> str: |
| """Pick a torch device: explicit preference, then CUDA, then CPU.""" |
| try: |
| import torch |
| except ImportError: |
| raise BrainError("model backends need torch: pip install 'tiny-computer-agent[vision]'") |
| if prefer: |
| return prefer |
| return "cuda" if torch.cuda.is_available() else "cpu" |
|
|
|
|
| def pick_dtype(device: str) -> Any: |
| import torch |
|
|
| if device.startswith("cuda"): |
| return torch.bfloat16 |
| return torch.float32 |
|
|
|
|
| def model_dependencies() -> tuple[bool, str]: |
| """(available, detail) for preflight and the doctor command.""" |
| try: |
| import transformers |
| except ImportError: |
| return False, "transformers is not installed (pip install 'tiny-computer-agent[vision]')" |
| try: |
| import torch |
| except ImportError: |
| return False, "torch is not installed (pip install 'tiny-computer-agent[vision]')" |
| return True, f"transformers {transformers.__version__}, torch {torch.__version__}" |
|
|
|
|
| def model_cached(model_id: str) -> bool: |
| """True when the weights are already in the local Hugging Face cache.""" |
| try: |
| from huggingface_hub import constants |
| except ImportError: |
| return False |
| folder = os.path.join(constants.HF_HUB_CACHE, "models--" + model_id.replace("/", "--")) |
| return os.path.isdir(folder) |
|
|
|
|
| class TextPlannerBrain(ModelBrain): |
| """Causal-LM planner over the OCR index: no pixels in, one menu number out.""" |
|
|
| name = "text-planner" |
| requires_index = True |
| supports_constrained = True |
|
|
| def __init__(self, model_id: str = DEFAULT_TEXT_MODEL, *, device: str | None = None, |
| **kwargs: Any) -> None: |
| super().__init__(model_id, **kwargs) |
| self.device = device |
| self._tokenizer: Any = None |
| self._model: Any = None |
|
|
| def _load(self) -> None: |
| if self._model is not None: |
| return |
| try: |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
| except ImportError as exc: |
| raise BrainError( |
| "text planner needs transformers and torch: " |
| "pip install 'tiny-computer-agent[vision]'" |
| ) from exc |
| self.device = resolve_device(self.device) |
| self._tokenizer = AutoTokenizer.from_pretrained(self.model_id) |
| self._model = AutoModelForCausalLM.from_pretrained( |
| self.model_id, dtype=pick_dtype(self.device) |
| ).to(self.device).eval() |
|
|
| def unload(self) -> None: |
| self._model = None |
| self._tokenizer = None |
|
|
| def _tokenizer_for_choice(self) -> Any: |
| self._load() |
| return self._tokenizer |
|
|
| def _generate_constrained(self, prompt: str, observation: Observation, |
| allowed_ids: Any, max_new_tokens: int, |
| *, sample: bool = False) -> str: |
| """One short generation that can only contain a menu number.""" |
| import torch |
|
|
| self._load() |
| messages = [{"role": "user", "content": prompt}] |
| text = self._tokenizer.apply_chat_template( |
| messages, add_generation_prompt=True, tokenize=False |
| ) |
| inputs = self._tokenizer(text, return_tensors="pt").to(self.device) |
| settings: dict[str, Any] = { |
| "max_new_tokens": max_new_tokens, |
| "pad_token_id": self._tokenizer.eos_token_id, |
| "logits_processor": _masked_generation( |
| self._model, allowed_ids, self._tokenizer.eos_token_id), |
| } |
| if sample: |
| settings.update(do_sample=True, temperature=0.7, top_p=0.9) |
| else: |
| settings["do_sample"] = False |
| with torch.inference_mode(): |
| output = self._model.generate(**inputs, **settings) |
| return self._tokenizer.decode( |
| output[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True |
| ).strip() |
|
|
| def _score_pair(self, prompt: str, observation: Observation, |
| yes_ids: Any, no_ids: Any) -> float: |
| import torch |
|
|
| self._load() |
| messages = [{"role": "user", "content": prompt}] |
| text = self._tokenizer.apply_chat_template( |
| messages, add_generation_prompt=True, tokenize=False) |
| inputs = self._tokenizer(text, return_tensors="pt").to(self.device) |
| with torch.inference_mode(): |
| out = self._model.generate( |
| **inputs, max_new_tokens=1, do_sample=False, return_dict_in_generate=True, |
| output_scores=True, pad_token_id=self._tokenizer.eos_token_id) |
| return _pair_margin(out.scores[0][0], yes_ids, no_ids) |
|
|
| def _generate(self, prompt: str, observation: Observation) -> str: |
| import torch |
|
|
| self._load() |
| messages = [{"role": "user", "content": prompt}] |
| text = self._tokenizer.apply_chat_template( |
| messages, add_generation_prompt=True, tokenize=False |
| ) |
| inputs = self._tokenizer(text, return_tensors="pt").to(self.device) |
| with torch.inference_mode(): |
| output = self._model.generate( |
| **inputs, |
| max_new_tokens=self.max_new_tokens, |
| do_sample=False, |
| pad_token_id=self._tokenizer.eos_token_id, |
| ) |
| return self._tokenizer.decode( |
| output[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True |
| ).strip() |
|
|
|
|
| class VLMBrain(ModelBrain): |
| """Tiny vision-language model: sees the screenshot, answers with a menu number.""" |
|
|
| name = "vlm" |
| requires_index = False |
| supports_constrained = True |
|
|
| def __init__(self, model_id: str = DEFAULT_VISION_MODEL, *, device: str | None = None, |
| **kwargs: Any) -> None: |
| super().__init__(model_id, **kwargs) |
| self.device = device |
| self._processor: Any = None |
| self._model: Any = None |
|
|
| def _load(self) -> None: |
| if self._model is not None: |
| return |
| try: |
| from transformers import AutoModelForImageTextToText, AutoProcessor |
| except ImportError as exc: |
| raise BrainError( |
| "the vision backend needs transformers and torch: " |
| "pip install 'tiny-computer-agent[vision]'" |
| ) from exc |
| self.device = resolve_device(self.device) |
| self._processor = AutoProcessor.from_pretrained(self.model_id) |
| self._model = AutoModelForImageTextToText.from_pretrained( |
| self.model_id, dtype=pick_dtype(self.device) |
| ).to(self.device).eval() |
|
|
| def unload(self) -> None: |
| self._model = None |
| self._processor = None |
|
|
| def _tokenizer_for_choice(self) -> Any: |
| self._load() |
| return self._processor.tokenizer |
|
|
| def _generate_constrained(self, prompt: str, observation: Observation, |
| allowed_ids: Any, max_new_tokens: int, |
| *, sample: bool = False) -> str: |
| """A vision generation that can only contain a menu number. |
| |
| This is the fix for the echo: the model is no longer able to reply with |
| the prompt, because the prompt's tokens are masked out of the logits. |
| """ |
| import torch |
|
|
| image = observation.frame |
| if image is None or isinstance(image, str): |
| raise BrainError( |
| "the vision backend needs a real screenshot but the observation carries " |
| f"{type(image).__name__} - run with --live, or use the text planner" |
| ) |
| self._load() |
| messages = [{"role": "user", "content": [{"type": "image"}, |
| {"type": "text", "text": prompt}]}] |
| text = self._processor.apply_chat_template(messages, add_generation_prompt=True) |
| inputs = self._processor(text=text, images=[image], return_tensors="pt").to(self.device) |
| settings: dict[str, Any] = { |
| "max_new_tokens": max_new_tokens, |
| "logits_processor": _masked_generation( |
| self._model, allowed_ids, |
| getattr(self._processor.tokenizer, "eos_token_id", None)), |
| } |
| if sample: |
| settings.update(do_sample=True, temperature=0.7, top_p=0.9) |
| else: |
| settings["do_sample"] = False |
| with torch.inference_mode(): |
| output = self._model.generate(**inputs, **settings) |
| return self._processor.batch_decode( |
| output[:, inputs["input_ids"].shape[1]:], skip_special_tokens=True |
| )[0].strip() |
|
|
| def _score_pair(self, prompt: str, observation: Observation, |
| yes_ids: Any, no_ids: Any) -> float: |
| """log P(yes) - log P(no) for one image question, read off the logits.""" |
| import torch |
|
|
| image = observation.frame |
| if image is None or isinstance(image, str): |
| raise BrainError( |
| "the vision backend needs a real screenshot but the observation carries " |
| f"{type(image).__name__} - run with --live, or use the text planner" |
| ) |
| self._load() |
| messages = [{"role": "user", "content": [{"type": "image"}, |
| {"type": "text", "text": prompt}]}] |
| text = self._processor.apply_chat_template(messages, add_generation_prompt=True) |
| inputs = self._processor(text=text, images=[image], return_tensors="pt").to(self.device) |
| with torch.inference_mode(): |
| out = self._model.generate(**inputs, max_new_tokens=1, do_sample=False, |
| return_dict_in_generate=True, output_scores=True) |
| return _pair_margin(out.scores[0][0], yes_ids, no_ids) |
|
|
| def _generate(self, prompt: str, observation: Observation) -> str: |
| import torch |
|
|
| image = observation.frame |
| if image is None or isinstance(image, str): |
| raise BrainError( |
| "the vision backend needs a real screenshot but the observation carries " |
| f"{type(image).__name__} - run with --live, or use the text planner" |
| ) |
| self._load() |
| messages = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": prompt}]}] |
| text = self._processor.apply_chat_template(messages, add_generation_prompt=True) |
| inputs = self._processor(text=text, images=[image], return_tensors="pt").to(self.device) |
| with torch.inference_mode(): |
| output = self._model.generate( |
| **inputs, max_new_tokens=self.max_new_tokens, do_sample=False |
| ) |
| return self._processor.batch_decode( |
| output[:, inputs["input_ids"].shape[1]:], skip_special_tokens=True |
| )[0].strip() |
|
|
|
|
| MODEL_BRAINS = {"text-planner": TextPlannerBrain, "vlm": VLMBrain} |
|
|
|
|
| def make_model_brain(name: str = "text-planner", **kwargs: Any) -> ModelBrain: |
| """Build a model brain by name, with an actionable error for unknown names.""" |
| if name not in MODEL_BRAINS: |
| raise BrainError(f"unknown brain {name!r}; expected one of {tuple(MODEL_BRAINS)}") |
| return MODEL_BRAINS[name](**kwargs) |
|
|
