| """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.observe import Observation |
|
|
| DEFAULT_TEXT_MODEL = "Qwen/Qwen2.5-0.5B-Instruct" |
| DEFAULT_VISION_MODEL = "HuggingFaceTB/SmolVLM-256M-Instruct" |
|
|
|
|
| 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 |
|
|
| 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 _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 |
|
|
| 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 _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) |
|
|
