openadapt-ml 0.1.0__py3-none-any.whl

openadapt_ml/models/qwen_vl.py

@@ -0,0 +1,426 @@
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+
+from PIL import Image
+import torch
+from peft import LoraConfig, PeftModel, get_peft_model
+from transformers import AutoProcessor, Qwen3VLForConditionalGeneration, Qwen2_5_VLForConditionalGeneration
+
+from openadapt_ml.models.base_adapter import BaseVLMAdapter, get_default_device
+
+
+def _process_vision_info(messages: List[Dict[str, Any]]) -> tuple[list[list[Any]], list[list[Any]]]:
+    """Minimal stand-in for qwen_vl_utils.process_vision_info.
+
+    For our use case we only need to extract image/video entries from the
+    message content structure expected by Qwen-VL examples, where each
+    message has a `content` list of dicts with `type` in {"image", "video"}.
+
+    Returns (image_inputs, video_inputs), each a list-of-lists suitable for
+    passing to AutoProcessor.
+    """
+
+    image_inputs: list[list[Any]] = []
+    video_inputs: list[list[Any]] = []
+
+    current_images: list[Any] = []
+    current_videos: list[Any] = []
+
+    for message in messages:
+        content = message.get("content", [])
+        if not isinstance(content, list):
+            continue
+        for item in content:
+            if not isinstance(item, dict):
+                continue
+            t = item.get("type")
+            if t == "image":
+                current_images.append(item.get("image"))
+            elif t == "video":
+                current_videos.append(item.get("video"))
+
+    if current_images:
+        image_inputs.append(current_images)
+    if current_videos:
+        video_inputs.append(current_videos)
+
+    return image_inputs, video_inputs
+
+
+class QwenVLAdapter(BaseVLMAdapter):
+    """Adapter for Qwen-family VLMs using Hugging Face + PEFT.
+
+    This is a minimal skeleton that:
+    - loads a base model + processor
+    - optionally applies a LoRA adapter
+    - implements the BaseVLMAdapter interface
+
+    The exact chat/image encoding and loss masking will be filled in
+    once we wire a concrete training loop.
+    """
+
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        processor: Any,
+        device: Optional[torch.device] = None,
+        version: str = "qwen3",
+    ) -> None:
+        super().__init__(model=model, processor=processor, device=device)
+        self.version = version
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_name: str,
+        lora_config: Optional[LoraConfig | Dict[str, Any]] = None,
+        load_in_4bit: bool = False,
+        device: Optional[torch.device] = None,
+        max_pixels: Optional[int] = None,
+        min_pixels: Optional[int] = None,
+    ) -> "QwenVLAdapter":
+        """Load base Qwen model + processor and attach optional LoRA adapter.
+
+        Args:
+            max_pixels: Maximum image size in pixels (e.g., 512*512=262144 for faster training).
+                        If None, uses model default (very large).
+            min_pixels: Minimum image size in pixels. If None, uses model default.
+        """
+
+        if "Qwen3-VL" in model_name or "Qwen3VL" in model_name:
+            version = "qwen3"
+            model_cls = Qwen3VLForConditionalGeneration
+        elif "Qwen2.5-VL" in model_name or "Qwen2_5" in model_name:
+            version = "qwen2_5"
+            model_cls = Qwen2_5_VLForConditionalGeneration
+        else:
+            raise ValueError(f"Unrecognized Qwen-VL model name: {model_name}")
+
+        processor = AutoProcessor.from_pretrained(model_name)
+
+        # Configure image resolution for faster training
+        if max_pixels is not None and hasattr(processor, 'image_processor'):
+            processor.image_processor.max_pixels = max_pixels
+            print(f"Set max_pixels to {max_pixels} ({int(max_pixels**0.5)}x{int(max_pixels**0.5)} approx)")
+        if min_pixels is not None and hasattr(processor, 'image_processor'):
+            processor.image_processor.min_pixels = min_pixels
+
+        model_kwargs: Dict[str, Any] = {}
+        if load_in_4bit:
+            model_kwargs["load_in_4bit"] = True
+
+        model = model_cls.from_pretrained(model_name, **model_kwargs)
+
+        # Support two modes for LoRA:
+        # - config-based (no weights_path): create a fresh adapter.
+        # - weights-based (weights_path in dict): load an existing adapter.
+        lora_weights_path: Optional[str] = None
+        lora_cfg_clean: Optional[LoraConfig | Dict[str, Any]] = None
+
+        if lora_config is not None:
+            if isinstance(lora_config, dict):
+                lora_weights_path = lora_config.get("weights_path")
+                lora_cfg_clean = {k: v for k, v in lora_config.items() if k != "weights_path"}
+            else:
+                lora_cfg_clean = lora_config
+
+        if lora_weights_path:
+            # Load an existing adapter onto the base model.
+            model = PeftModel.from_pretrained(model, lora_weights_path)
+        elif lora_cfg_clean is not None:
+            if isinstance(lora_cfg_clean, dict):
+                lora_cfg_clean = LoraConfig(**lora_cfg_clean)
+            model = get_peft_model(model, lora_cfg_clean)  # type: ignore[arg-type]
+
+        if device is None:
+            device = get_default_device()
+
+        return cls(model=model, processor=processor, device=device, version=version)
+
+    def prepare_inputs(self, batch: List[Dict[str, Any]]) -> Dict[str, Any]:  # type: ignore[override]
+        """Convert SFT samples into model inputs for Qwen-VL.
+
+        Supports true batching by processing multiple samples simultaneously.
+        Uses processor.apply_chat_template with padding=True and truncation=True
+        for multi-sample tokenization. Computes assistant-only labels per sample.
+
+        Args:
+            batch: List of SFT-style samples, each with "images" and "messages" keys.
+
+        Returns:
+            Dict with input_ids, attention_mask, pixel_values, image_grid_thw, and labels.
+        """
+
+        if self.version == "qwen3":
+            # Build batch of messages for all samples
+            batch_messages_full: List[List[Dict[str, Any]]] = []
+            batch_messages_user_only: List[List[Dict[str, Any]]] = []
+
+            for sample in batch:
+                image_paths = sample["images"]
+                if not image_paths:
+                    raise ValueError("Sample is missing image paths")
+                image_path = image_paths[0]
+
+                messages = sample["messages"]
+                user_text = ""
+                assistant_text = ""
+                for m in messages:
+                    role = m.get("role")
+                    if role == "user":
+                        user_text = m.get("content", "")
+                    elif role == "assistant":
+                        assistant_text = m.get("content", "")
+
+                # Full messages (user + assistant)
+                qwen_messages_full: List[Dict[str, Any]] = [
+                    {
+                        "role": "user",
+                        "content": [
+                            {"type": "image", "image": image_path},
+                            {"type": "text", "text": user_text},
+                        ],
+                    },
+                ]
+                if assistant_text:
+                    qwen_messages_full.append({
+                        "role": "assistant",
+                        "content": [{"type": "text", "text": assistant_text}],
+                    })
+                batch_messages_full.append(qwen_messages_full)
+
+                # User-only messages (for label masking)
+                qwen_messages_user_only: List[Dict[str, Any]] = [
+                    {
+                        "role": "user",
+                        "content": [
+                            {"type": "image", "image": image_path},
+                            {"type": "text", "text": user_text},
+                        ],
+                    }
+                ]
+                batch_messages_user_only.append(qwen_messages_user_only)
+
+            # Tokenize full batch with padding and truncation
+            inputs_full = self.processor.apply_chat_template(  # type: ignore[call-arg]
+                batch_messages_full,
+                tokenize=True,
+                add_generation_prompt=False,
+                return_dict=True,
+                return_tensors="pt",
+                padding=True,
+                truncation=True,
+            )
+
+            # Tokenize user-only batch for label masking
+            inputs_user = self.processor.apply_chat_template(  # type: ignore[call-arg]
+                batch_messages_user_only,
+                tokenize=True,
+                add_generation_prompt=True,
+                return_dict=True,
+                return_tensors="pt",
+                padding=True,
+                truncation=True,
+            )
+
+            input_ids_full = inputs_full["input_ids"]  # [batch_size, seq_len]
+            input_ids_user = inputs_user["input_ids"]  # [batch_size, seq_len_user]
+
+            # Initialize labels with full input_ids, then mask per sample
+            labels = input_ids_full.clone()
+
+            # Compute assistant-only labels per sample
+            batch_size = input_ids_full.size(0)
+            for i in range(batch_size):
+                sample = batch[i]
+                messages = sample["messages"]
+                assistant_text = ""
+                for m in messages:
+                    if m.get("role") == "assistant":
+                        assistant_text = m.get("content", "")
+
+                if assistant_text:
+                    # Find where user prompt ends and assistant response begins
+                    # The user-only sequence should be a prefix of the full sequence
+                    full_ids = input_ids_full[i]
+                    user_ids = input_ids_user[i]
+
+                    # Remove padding from user_ids to find actual sequence length
+                    # Padding token is typically 0 or a special value
+                    # For Qwen, we look for the first occurrence of pad token
+                    pad_token_id = self.processor.tokenizer.pad_token_id
+                    user_ids_no_pad = user_ids[user_ids != pad_token_id] if pad_token_id is not None else user_ids
+                    user_len = len(user_ids_no_pad)
+
+                    # Check if user sequence is a prefix of full sequence
+                    if user_len <= full_ids.size(0):
+                        # Mask everything except assistant tokens
+                        labels[i, :] = -100
+                        # Only supervise on tokens after the user prompt
+                        labels[i, user_len:] = full_ids[user_len:]
+
+            # Ensure padding tokens are masked in labels
+            if hasattr(self.processor.tokenizer, 'pad_token_id') and self.processor.tokenizer.pad_token_id is not None:
+                labels[input_ids_full == self.processor.tokenizer.pad_token_id] = -100
+
+            inputs_full["labels"] = labels
+            return inputs_full
+
+        else:  # qwen2_5
+            # Build batch of messages
+            batch_messages: List[List[Dict[str, Any]]] = []
+            batch_texts: List[str] = []
+            all_image_inputs: List[List[Any]] = []
+            all_video_inputs: List[List[Any]] = []
+
+            for sample in batch:
+                image_paths = sample["images"]
+                if not image_paths:
+                    raise ValueError("Sample is missing image paths")
+                image_path = image_paths[0]
+
+                messages = sample["messages"]
+                user_text = ""
+                assistant_text = ""
+                for m in messages:
+                    role = m.get("role")
+                    if role == "user":
+                        user_text = m.get("content", "")
+                    elif role == "assistant":
+                        assistant_text = m.get("content", "")
+
+                qwen_messages: List[Dict[str, Any]] = [
+                    {
+                        "role": "user",
+                        "content": [
+                            {"type": "image", "image": image_path},
+                            {"type": "text", "text": user_text},
+                        ],
+                    }
+                ]
+                if assistant_text:
+                    qwen_messages.append({
+                        "role": "assistant",
+                        "content": [{"type": "text", "text": assistant_text}],
+                    })
+
+                batch_messages.append(qwen_messages)
+
+                # Convert to text for this sample
+                text = self.processor.apply_chat_template(  # type: ignore[call-arg]
+                    qwen_messages,
+                    tokenize=False,
+                    add_generation_prompt=False,
+                )
+                batch_texts.append(text)
+
+                # Extract vision info
+                image_inputs, video_inputs = _process_vision_info(qwen_messages)
+                if image_inputs:
+                    all_image_inputs.extend(image_inputs)
+                if video_inputs:
+                    all_video_inputs.extend(video_inputs)
+
+            # Process batch with padding
+            videos_arg = all_video_inputs if all_video_inputs else None
+            inputs = self.processor(  # type: ignore[call-arg]
+                text=batch_texts,
+                images=all_image_inputs if all_image_inputs else None,
+                videos=videos_arg,
+                padding=True,
+                truncation=True,
+                return_tensors="pt",
+            )
+
+            # For qwen2_5, use full sequence supervision for now
+            # (can be refined to assistant-only if needed)
+            input_ids = inputs["input_ids"]
+            labels = input_ids.clone()
+
+            # Mask padding tokens
+            if hasattr(self.processor.tokenizer, 'pad_token_id') and self.processor.tokenizer.pad_token_id is not None:
+                labels[input_ids == self.processor.tokenizer.pad_token_id] = -100
+
+            inputs["labels"] = labels
+            return inputs
+
+    def compute_loss(self, inputs: Dict[str, Any]) -> torch.Tensor:  # type: ignore[override]
+        inputs = {k: v.to(self.device) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}
+        outputs = self.model(**inputs)
+        # Hugging Face causal LM models return `loss` when `labels` are provided.
+        return outputs.loss  # type: ignore[no-any-return]
+
+    def generate(self, sample: Dict[str, Any], max_new_tokens: int = 64) -> str:  # type: ignore[override]
+        """Generate assistant text for a single SFT-style sample.
+
+        We pass system + user messages to the chat template with
+        `add_generation_prompt=True` and let the model generate the
+        assistant continuation.
+        """
+
+        image_paths = sample["images"]
+        if not image_paths:
+            raise ValueError("Sample is missing image paths")
+        image_path = image_paths[0]
+
+        messages = sample["messages"]
+        user_text = ""
+        for m in messages:
+            if m.get("role") == "user":
+                user_text = m.get("content", "")
+
+        qwen_messages: List[Dict[str, Any]] = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "image", "image": image_path},
+                    {"type": "text", "text": user_text},
+                ],
+            }
+        ]
+
+        if self.version == "qwen3":
+            inputs = self.processor.apply_chat_template(  # type: ignore[call-arg]
+                qwen_messages,
+                tokenize=True,
+                add_generation_prompt=True,
+                return_dict=True,
+                return_tensors="pt",
+            ).to(self.device)
+        else:
+            text = self.processor.apply_chat_template(  # type: ignore[call-arg]
+                qwen_messages,
+                tokenize=False,
+                add_generation_prompt=True,
+            )
+            image_inputs, video_inputs = _process_vision_info(qwen_messages)
+            videos_arg = video_inputs if video_inputs else None
+            inputs = self.processor(  # type: ignore[call-arg]
+                text=[text],
+                images=image_inputs,
+                videos=videos_arg,
+                padding=True,
+                return_tensors="pt",
+            ).to(self.device)
+
+        with torch.no_grad():
+            generation = self.model.generate(**inputs, max_new_tokens=max_new_tokens)
+
+        # Decode only the GENERATED tokens (not the input prompt)
+        # This is critical - otherwise we return "user Goal: ... assistant ..." instead of just the response
+        input_len = inputs["input_ids"].shape[1]
+        generated_ids = generation[:, input_len:]
+        text = self.processor.batch_decode(  # type: ignore[call-arg]
+            generated_ids,
+            skip_special_tokens=True,
+        )[0]
+        return text
+
+    def save_checkpoint(self, path: str) -> None:
+        """Save the LoRA adapter weights to a directory."""
+        from pathlib import Path
+        save_path = Path(path)
+        save_path.mkdir(parents=True, exist_ok=True)
+        # Save the PEFT adapter (LoRA weights only, not base model)
+        self.model.save_pretrained(str(save_path))

openadapt_ml/runtime/policy.py

@@ -0,0 +1,182 @@
+from __future__ import annotations
+
+import json
+import re
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple
+
+from PIL import Image
+
+from openadapt_ml.models.base_adapter import BaseVLMAdapter
+from openadapt_ml.schemas.sessions import Action
+
+
+# Coordinate-based DSL patterns
+_CLICK_RE = re.compile(r"CLICK\(x=([0-9]*\.?[0-9]+),\s*y=([0-9]*\.?[0-9]+)\)")
+_TYPE_RE = re.compile(r'TYPE\(text="([^"\\]*(?:\\.[^"\\]*)*)"\)')
+_WAIT_RE = re.compile(r"\bWAIT\s*\(\s*\)")
+_DONE_RE = re.compile(r"\bDONE\s*\(\s*\)")
+
+# SoM (Set-of-Marks) index-based DSL patterns
+_CLICK_SOM_RE = re.compile(r"CLICK\(\[(\d+)\]\)")
+_TYPE_SOM_RE = re.compile(r'TYPE\(\[(\d+)\],\s*["\']([^"\']*(?:\\.[^"\']*)*)["\']\)')
+_TYPE_SOM_SIMPLE_RE = re.compile(r'TYPE\(["\']([^"\']*(?:\\.[^"\']*)*)["\']\)')
+
+
+@dataclass
+class PolicyOutput:
+    """Result of a single policy step."""
+    action: Action
+    thought: Optional[str] = None
+    state: Optional[Dict[str, Any]] = None
+    raw_text: str = ""
+
+
+def parse_thought_state_action(text: str) -> Tuple[Optional[str], Optional[Dict[str, Any]], str]:
+    """Parse Thought / State / Action blocks from model output.
+
+    Expected format:
+        Thought: [reasoning]
+        State: {"success": false, "progress": 0.5, ...}
+        Action: CLICK(x=0.42, y=0.31)
+
+    Returns:
+        (thought, state, action_str):
+        - thought: Content after 'Thought:' up to 'State:' or 'Action:'
+        - state: Parsed JSON dict from 'State:' line, or None
+        - action_str: Content after 'Action:', or whole text if missing
+
+    Note: We look for the LAST occurrence of 'Action:' to handle cases where
+    the user prompt template also contains 'Action:' placeholders.
+    """
+    thought: Optional[str] = None
+    state: Optional[Dict[str, Any]] = None
+    action_str: str = text.strip()
+
+    # Extract Thought - find the LAST occurrence (model's response, not template)
+    thought_matches = list(re.finditer(r"Thought:\s*(.+?)(?=State:|Action:|$)", text, re.DOTALL | re.IGNORECASE))
+    if thought_matches:
+        thought = thought_matches[-1].group(1).strip()
+
+    # Extract State (JSON on same line or next line) - last occurrence
+    state_matches = list(re.finditer(r"State:\s*(\{.*?\})", text, re.DOTALL | re.IGNORECASE))
+    if state_matches:
+        try:
+            state = json.loads(state_matches[-1].group(1))
+        except json.JSONDecodeError:
+            state = None
+
+    # Extract Action - find the LAST occurrence to get the model's actual action
+    # not the placeholder in the prompt template
+    action_matches = list(re.finditer(r"Action:\s*(.+?)(?=\n|$)", text, re.IGNORECASE))
+    if action_matches:
+        action_str = action_matches[-1].group(1).strip()
+
+    return thought, state, action_str
+
+
+class AgentPolicy:
+    """Runtime policy wrapper around a trained VLM adapter.
+
+    Formats goal-conditioned inputs and parses textual actions into
+    structured `Action` objects.
+    """
+
+    def __init__(self, adapter: BaseVLMAdapter) -> None:
+        self.adapter = adapter
+
+    def _build_sample(self, image: Image.Image, goal: str) -> Dict[str, Any]:
+        # For runtime we keep the same structure as SFT samples but use
+        # an in-memory image. The adapter's generate method currently expects
+        # paths, so we require the caller to supply a path-based sample. For
+        # now, we save responsibility for image loading to the caller; this
+        # method is kept for future extensibility.
+        raise NotImplementedError(
+            "AgentPolicy._build_sample is not used directly; pass a sample dict "
+            "compatible with the adapter's `generate` method."
+        )
+
+    def _parse_action(self, text: str) -> Action:
+        """Parse a DSL action string into an Action object.
+
+        Supported formats (coordinate-based):
+        - CLICK(x=<float>, y=<float>)
+        - TYPE(text="...")
+
+        Supported formats (SoM index-based):
+        - CLICK([N])
+        - TYPE([N], "text")
+        - TYPE("text")
+
+        Common formats:
+        - WAIT()
+        - DONE()
+
+        Returns Action(type="failed") if no valid action is found.
+        """
+        # Try SoM patterns first (index-based)
+        # CLICK([N])
+        m = _CLICK_SOM_RE.search(text)
+        if m:
+            idx = int(m.group(1))
+            return Action(type="click", element_index=idx)
+
+        # TYPE([N], "text")
+        m = _TYPE_SOM_RE.search(text)
+        if m:
+            idx = int(m.group(1))
+            raw_text = m.group(2)
+            unescaped = raw_text.replace('\\"', '"').replace("\\\\", "\\")
+            return Action(type="type", text=unescaped, element_index=idx)
+
+        # TYPE("text") - SoM style without index
+        m = _TYPE_SOM_SIMPLE_RE.search(text)
+        if m:
+            raw_text = m.group(1)
+            unescaped = raw_text.replace('\\"', '"').replace("\\\\", "\\")
+            return Action(type="type", text=unescaped)
+
+        # Coordinate-based patterns
+        # CLICK(x=..., y=...)
+        m = _CLICK_RE.search(text)
+        if m:
+            x = float(m.group(1))
+            y = float(m.group(2))
+            # Clamp to [0, 1]
+            x = max(0.0, min(1.0, x))
+            y = max(0.0, min(1.0, y))
+            return Action(type="click", x=x, y=y)
+
+        # TYPE(text="...")
+        m = _TYPE_RE.search(text)
+        if m:
+            # Unescape the text content
+            raw_text = m.group(1)
+            unescaped = raw_text.replace('\\"', '"').replace("\\\\", "\\")
+            return Action(type="type", text=unescaped)
+
+        # WAIT()
+        if _WAIT_RE.search(text):
+            return Action(type="wait")
+
+        # DONE()
+        if _DONE_RE.search(text):
+            return Action(type="done")
+
+        # Fallback
+        return Action(type="failed", raw={"text": text})
+
+    def predict_action_from_sample(
+        self, sample: Dict[str, Any], max_new_tokens: int = 150
+    ) -> Tuple[Action, Optional[str], Optional[Dict[str, Any]], str]:
+        """Run the adapter on a pre-built SFT-style sample and parse the result.
+
+        Returns (Action, thought, state, raw_text) where:
+        - thought: Reasoning text from 'Thought:' block
+        - state: Parsed JSON dict from 'State:' block (may contain 'success' bool)
+        - raw_text: The raw model output text for debugging
+        """
+        text = self.adapter.generate(sample, max_new_tokens=max_new_tokens)
+        thought, state, action_str = parse_thought_state_action(text)
+        action = self._parse_action(action_str)
+        return action, thought, state, text

openadapt_ml/schemas/__init__.py

@@ -0,0 +1,53 @@
+"""Schema definitions and validation for openadapt-ml.
+
+Core data structures:
+    - Action: A single GUI action (click, type, scroll, etc.)
+    - Observation: GUI state observation (screenshot, accessibility tree, etc.)
+    - Step: One timestep containing observation + action
+    - Episode: A single task attempt / workflow instance
+    - Session: Container for multiple episodes
+
+Validation:
+    - validate_episode(): Validate an Episode object
+    - validate_session(): Validate a Session object
+    - validate_episodes(): Validate a list of Episodes
+    - ValidationError: Raised on schema violations
+"""
+
+from openadapt_ml.schemas.sessions import (
+    Action,
+    ActionType,
+    Episode,
+    Observation,
+    Session,
+    Step,
+)
+from openadapt_ml.schemas.validation import (
+    ValidationError,
+    summarize_episodes,
+    validate_action,
+    validate_episode,
+    validate_episodes,
+    validate_observation,
+    validate_session,
+    validate_step,
+)
+
+__all__ = [
+    # Core types
+    "Action",
+    "ActionType",
+    "Episode",
+    "Observation",
+    "Session",
+    "Step",
+    # Validation
+    "ValidationError",
+    "validate_action",
+    "validate_episode",
+    "validate_episodes",
+    "validate_observation",
+    "validate_session",
+    "validate_step",
+    "summarize_episodes",
+]