|
| 1 | +"""Qwen3-VL backend using HuggingFace transformers. |
| 2 | +
|
| 3 | +Qwen3-VL (Alibaba, Sept 2025) is an open vision-language model with a |
| 4 | +ViT vision encoder and accessible cross-modal attention. Unlike CLIP β |
| 5 | +a dual-encoder that only scores image/text similarity β Qwen3-VL is a |
| 6 | +generative VLM, so its attention reflects genuine cross-modal reasoning |
| 7 | +between the question tokens and the image patches. That makes it the |
| 8 | +first miru backend whose saliency the synergy probe (``miru.synergy``) |
| 9 | +and the deletion test (``miru.fidelity``) can meaningfully interrogate. |
| 10 | +
|
| 11 | +Attention extraction |
| 12 | +-------------------- |
| 13 | +A forward pass with ``output_attentions=True`` yields per-layer maps of |
| 14 | +shape ``(batch, heads, seq, seq)``. We: |
| 15 | +
|
| 16 | +1. pick a **middle** decoder layer β cross-modal fusion concentrates in |
| 17 | + the middle of the stack (Qwen2.5-VL technical report, arXiv:2502.13923, |
| 18 | + found layers ~14-24 carry the peak vision-language signal); |
| 19 | +2. read the **last prompt token's** attention to the **image tokens** |
| 20 | + (located via ``config.image_token_id``), averaged over heads; |
| 21 | +3. reshape that 1-D vector to the nearest square patch grid. |
| 22 | +
|
| 23 | +The numeric reshaping logic lives in pure helpers |
| 24 | +(:func:`_select_middle_layer`, :func:`_attention_row_to_grid`) so it is |
| 25 | +unit-tested offline; the model load + generation path is exercised only |
| 26 | +under ``MIRU_TEST_REAL_BACKENDS=1``. |
| 27 | +
|
| 28 | +Requires: ``pip install 'transformers>=4.57' torch`` (Qwen3-VL is |
| 29 | +integrated natively in recent transformers). ``output_attentions`` |
| 30 | +requires the eager attention implementation. |
| 31 | +""" |
| 32 | + |
| 33 | +from __future__ import annotations |
| 34 | + |
| 35 | +import numpy as np |
| 36 | + |
| 37 | +from miru.models.base import VLMBackend, VLMOutput |
| 38 | + |
| 39 | +_DEFAULT_MODEL = "Qwen/Qwen3-VL-8B-Instruct" |
| 40 | +_MIDDLE_LAYER_FRACTION = 0.6 # ~60% deep β within the empirical fusion band. |
| 41 | + |
| 42 | + |
| 43 | +def _select_middle_layer( |
| 44 | + num_layers: int, fraction: float = _MIDDLE_LAYER_FRACTION |
| 45 | +) -> int: |
| 46 | + """Index of the decoder layer ``fraction`` of the way up the stack. |
| 47 | +
|
| 48 | + Cross-modal fusion peaks mid-stack, so we sample there rather than at |
| 49 | + the final layer. Clamped to a valid ``[0, num_layers - 1]`` index. |
| 50 | + """ |
| 51 | + if num_layers < 1: |
| 52 | + raise ValueError(f"num_layers must be >= 1, got {num_layers}") |
| 53 | + idx = int(round((num_layers - 1) * fraction)) |
| 54 | + return max(0, min(num_layers - 1, idx)) |
| 55 | + |
| 56 | + |
| 57 | +def _attention_row_to_grid(attn_row: np.ndarray) -> np.ndarray: |
| 58 | + """Reshape a 1-D image-token attention vector to a square float32 grid. |
| 59 | +
|
| 60 | + The vector is truncated to the largest square ``g*g <= len`` and |
| 61 | + reshaped to ``(g, g)`` β mirroring the CLIP backend's patch-grid |
| 62 | + derivation. |
| 63 | + """ |
| 64 | + flat = np.asarray(attn_row, dtype=np.float32).ravel() |
| 65 | + if flat.size == 0: |
| 66 | + raise ValueError("attn_row must be non-empty") |
| 67 | + grid = int(flat.size**0.5) |
| 68 | + if grid < 1: |
| 69 | + raise ValueError(f"too few image tokens to form a grid: {flat.size}") |
| 70 | + return flat[: grid * grid].reshape(grid, grid).astype(np.float32) |
| 71 | + |
| 72 | + |
| 73 | +class Qwen3VLBackend(VLMBackend): |
| 74 | + """Generative VLM backend over Qwen3-VL with cross-modal attention. |
| 75 | +
|
| 76 | + Lazy-loads weights on the first :meth:`infer` call β never at import |
| 77 | + time β so the module imports cleanly in mock-only environments. |
| 78 | + """ |
| 79 | + |
| 80 | + def __init__(self, model_name: str = _DEFAULT_MODEL) -> None: |
| 81 | + self._model_name = model_name |
| 82 | + self._model = None |
| 83 | + self._processor = None |
| 84 | + |
| 85 | + def _load(self) -> None: |
| 86 | + """Lazy-load model + processor on first inference call.""" |
| 87 | + if self._model is not None: |
| 88 | + return |
| 89 | + import torch # noqa: F401 -- kept out of module scope (optional dep) |
| 90 | + from transformers import AutoProcessor, Qwen3VLForConditionalGeneration |
| 91 | + |
| 92 | + self._processor = AutoProcessor.from_pretrained(self._model_name) |
| 93 | + self._model = Qwen3VLForConditionalGeneration.from_pretrained( |
| 94 | + self._model_name, |
| 95 | + attn_implementation="eager", # required for output_attentions |
| 96 | + ) |
| 97 | + self._model.eval() |
| 98 | + |
| 99 | + @property |
| 100 | + def name(self) -> str: |
| 101 | + return "qwen3vl" |
| 102 | + |
| 103 | + def infer(self, image_array: np.ndarray, question: str) -> VLMOutput: |
| 104 | + """Run Qwen3-VL inference. |
| 105 | +
|
| 106 | + Args: |
| 107 | + image_array: float32 ``(H, W, 3)`` in ``[0, 1]`` β converted to |
| 108 | + a uint8 PIL image internally. |
| 109 | + question: The user question conditioning the model. |
| 110 | +
|
| 111 | + Returns: |
| 112 | + ``VLMOutput`` with the generated answer, a confidence derived |
| 113 | + from the first generated token's probability, a square |
| 114 | + cross-modal attention grid, and reasoning steps. |
| 115 | + """ |
| 116 | + import torch |
| 117 | + from PIL import Image |
| 118 | + |
| 119 | + self._load() |
| 120 | + |
| 121 | + img_uint8 = (np.clip(image_array, 0, 1) * 255).astype(np.uint8) |
| 122 | + pil_image = Image.fromarray(img_uint8) |
| 123 | + |
| 124 | + messages = [ |
| 125 | + { |
| 126 | + "role": "user", |
| 127 | + "content": [ |
| 128 | + {"type": "image", "image": pil_image}, |
| 129 | + {"type": "text", "text": question}, |
| 130 | + ], |
| 131 | + } |
| 132 | + ] |
| 133 | + prompt = self._processor.apply_chat_template( |
| 134 | + messages, tokenize=False, add_generation_prompt=True |
| 135 | + ) |
| 136 | + inputs = self._processor(text=[prompt], images=[pil_image], return_tensors="pt") |
| 137 | + |
| 138 | + with torch.no_grad(): |
| 139 | + generated = self._model.generate( |
| 140 | + **inputs, |
| 141 | + max_new_tokens=64, |
| 142 | + output_scores=True, |
| 143 | + return_dict_in_generate=True, |
| 144 | + ) |
| 145 | + forward = self._model(**inputs, output_attentions=True) |
| 146 | + |
| 147 | + answer = self._decode_answer(inputs, generated) |
| 148 | + confidence = _first_token_confidence(generated) |
| 149 | + attention_map = self._extract_attention(inputs, forward) |
| 150 | + |
| 151 | + return VLMOutput( |
| 152 | + answer=answer, |
| 153 | + confidence=confidence, |
| 154 | + attention_weights=attention_map, |
| 155 | + reasoning_steps=[ |
| 156 | + f"Encoded image + question with Qwen3-VL ({self._model_name})", |
| 157 | + f"Generated answer: '{answer[:50]}'", |
| 158 | + "Read cross-modal attention from a middle decoder layer", |
| 159 | + ], |
| 160 | + ) |
| 161 | + |
| 162 | + def _decode_answer(self, inputs: object, generated: object) -> str: |
| 163 | + """Strip the prompt tokens and decode only the generated continuation.""" |
| 164 | + input_len = inputs["input_ids"].shape[1] # type: ignore[index] |
| 165 | + new_tokens = generated.sequences[0][input_len:] # type: ignore[attr-defined] |
| 166 | + return self._processor.batch_decode( # type: ignore[union-attr] |
| 167 | + [new_tokens], skip_special_tokens=True |
| 168 | + )[0].strip() |
| 169 | + |
| 170 | + def _extract_attention(self, inputs: object, forward: object) -> np.ndarray: |
| 171 | + """Middle-layer, last-token attention over image tokens β square grid.""" |
| 172 | + attentions = forward.attentions # type: ignore[attr-defined] |
| 173 | + layer = _select_middle_layer(len(attentions)) |
| 174 | + layer_attn = attentions[layer][0] # (heads, seq, seq) |
| 175 | + image_token_id = self._model.config.image_token_id # type: ignore[union-attr] |
| 176 | + ids = inputs["input_ids"][0] # type: ignore[index] |
| 177 | + image_positions = (ids == image_token_id).nonzero(as_tuple=True)[0] |
| 178 | + # Last prompt token's head-averaged attention to the image tokens. |
| 179 | + last_to_image = layer_attn[:, -1, :].mean(dim=0)[image_positions] |
| 180 | + return _attention_row_to_grid(last_to_image.float().cpu().numpy()) |
| 181 | + |
| 182 | + |
| 183 | +def _first_token_confidence(generated: object) -> float: |
| 184 | + """Confidence = softmax probability of the first generated token. |
| 185 | +
|
| 186 | + A real signal of the model's certainty in its answer, clamped to |
| 187 | + ``[0, 1]``. Falls back to ``0.5`` when scores are unavailable. |
| 188 | + """ |
| 189 | + import torch |
| 190 | + |
| 191 | + scores = getattr(generated, "scores", None) |
| 192 | + if not scores: |
| 193 | + return 0.5 |
| 194 | + probs = torch.softmax(scores[0][0], dim=-1) |
| 195 | + return float(np.clip(float(probs.max().item()), 0.0, 1.0)) |
| 196 | + |
| 197 | + |
| 198 | +__all__ = [ |
| 199 | + "Qwen3VLBackend", |
| 200 | +] |
0 commit comments