| """The decision model around a Hugging Face causal LM: two-way reading and a code-row readout. |
| |
| **Two-way reading.** With ``attention="two_way"`` the full-attention layers see the whole prompt |
| in both directions, while the linear-attention layers (Qwen3.5's Gated DeltaNet) and any |
| sliding-window layers keep the model's own masks. |
| |
| **Readout.** The hidden state at each question's readout position is scored against one row per |
| option code: the LM head's own rows of the code tokens (``readout="lm_head"``), or a separate |
| ``[n_codes, d]`` matrix (``readout="head"``). The logits are normalised over the question's |
| options only. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| from typing import Any, Sequence |
|
|
| import torch |
| from torch import nn |
|
|
| from glance_runtime.layout import Example |
|
|
|
|
| def _core(lm: nn.Module) -> nn.Module: |
| """The decoder stack (the module that builds attention masks); PEFT wrappers are unwrapped.""" |
| if hasattr(lm, "get_base_model"): |
| lm = lm.get_base_model() |
| for attr in ("model", "language_model", "transformer"): |
| m = getattr(lm, attr, None) |
| if isinstance(m, nn.Module): |
| inner = getattr(m, "language_model", None) |
| return inner if isinstance(inner, nn.Module) else m |
| return lm |
|
|
|
|
| def _multimodal(lm: nn.Module) -> nn.Module | None: |
| """The module holding the vision encoder and M-RoPE indexing (Qwen3.5's ``model``), if any.""" |
| if hasattr(lm, "get_base_model"): |
| lm = lm.get_base_model() |
| for m in (getattr(lm, "model", None), lm): |
| if isinstance(m, nn.Module) and hasattr(m, "get_image_features") and hasattr(m, "get_rope_index") \ |
| and hasattr(m, "get_placeholder_mask"): |
| return m |
| return None |
|
|
|
|
| def _text_config(lm: nn.Module): |
| cfg = lm.config |
| return getattr(cfg, "text_config", None) or cfg |
|
|
|
|
| def two_way_masks(lm: nn.Module, inputs_embeds: torch.Tensor, attention_mask: torch.Tensor) -> Any: |
| """Bidirectional masks for full-attention layers; the model's own masks for the others.""" |
| from transformers import masking_utils as mu |
|
|
| cfg = _core(lm).config |
| kw = {"config": cfg, "inputs_embeds": inputs_embeds, "attention_mask": attention_mask, "past_key_values": None} |
| full = mu.create_bidirectional_mask(**kw, allow_is_bidirectional_skip=False) |
| layer_types = getattr(cfg, "layer_types", None) |
| if not layer_types: |
| return full |
| masks: dict[str, Any] = {} |
| pos = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device)[None].expand(inputs_embeds.shape[0], -1) |
| for t in set(layer_types): |
| if t == "full_attention": |
| masks[t] = full |
| elif t == "linear_attention": |
| masks[t] = mu.create_recurrent_attention_mask(**kw) |
| elif t == "sliding_attention": |
| masks[t] = mu.create_sliding_window_causal_mask(**kw, position_ids=pos) |
| else: |
| masks[t] = mu.create_causal_mask(**kw, position_ids=pos) |
| return masks |
|
|
|
|
| def varlen_masks(lm: nn.Module, inputs_embeds: torch.Tensor, attention_mask: torch.Tensor, causal: bool) -> dict: |
| """Masks for ``attn_impl="varlen"``: packing indices for full-attention layers (two-way or causal), |
| the model's own masks for the others.""" |
| from transformers import masking_utils as mu |
|
|
| from glance_runtime.attention import VarlenMask |
|
|
| cfg = _core(lm).config |
| kw = {"config": cfg, "inputs_embeds": inputs_embeds, "attention_mask": attention_mask, "past_key_values": None} |
| pos = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device)[None].expand(inputs_embeds.shape[0], -1) |
| masks: dict[str, Any] = {} |
| for t in set(getattr(cfg, "layer_types", None) or ["full_attention"]): |
| if t == "full_attention": |
| masks[t] = VarlenMask.from_padding(attention_mask, causal=causal) |
| elif t == "linear_attention": |
| masks[t] = mu.create_recurrent_attention_mask(**kw) |
| else: |
| raise ValueError(f"attn_impl varlen does not handle {t!r} layers") |
| return masks |
|
|
|
|
| def set_full_attention_causal(lm: nn.Module, causal: bool) -> int: |
| """Flip ``is_causal`` on full-attention modules (SDPA and flash paths read it). Returns how many.""" |
| n = 0 |
| for name, mod in lm.named_modules(): |
| cls = type(mod).__name__ |
| if cls.endswith("Attention") and "Vision" not in cls and hasattr(mod, "is_causal") and \ |
| getattr(mod, "layer_type", "full_attention") != "sliding_attention": |
| mod.is_causal = causal |
| n += 1 |
| return n |
|
|
|
|
| @dataclass |
| class Batch: |
| input_ids: torch.Tensor |
| attention_mask: torch.Tensor |
| readout_row: torch.Tensor |
| readout_pos: torch.Tensor |
| option_code: torch.Tensor |
| option_readout: torch.Tensor |
| sizes: list[int] |
| readouts: list |
| pixel_values: torch.Tensor | None = None |
| image_grid_thw: torch.Tensor | None = None |
|
|
| def to(self, device: Any) -> "Batch": |
| """Move the tensors the forward pass reads (indices stay where ``forward`` needs them).""" |
| self.input_ids = self.input_ids.to(device) |
| self.attention_mask = self.attention_mask.to(device) |
| if self.pixel_values is not None: |
| self.pixel_values = self.pixel_values.to(device) |
| self.image_grid_thw = self.image_grid_thw.to(device) |
| return self |
|
|
|
|
| def collate(examples: Sequence[Example], code_index: dict[str, int], pad_id: int = 0) -> Batch: |
| t = max(len(e.input_ids) for e in examples) |
| ids = torch.full((len(examples), t), pad_id, dtype=torch.long) |
| mask = torch.zeros((len(examples), t), dtype=torch.long) |
| rows, poss, codes, owners, sizes, ros = [], [], [], [], [], [] |
| for b, e in enumerate(examples): |
| ids[b, : len(e.input_ids)] = torch.tensor(e.input_ids) |
| mask[b, : len(e.input_ids)] = 1 |
| for ro in e.readouts: |
| r = len(ros) |
| rows.append(b) |
| poss.append(ro.position) |
| codes.extend(code_index[c] for c in ro.codes) |
| owners.extend([r] * len(ro.codes)) |
| sizes.append(len(ro.codes)) |
| ros.append(ro) |
| vision = [v for e in examples for v in e.vision] |
| pixels = torch.cat([v["pixel_values"] for v in vision]) if vision else None |
| grids = torch.cat([v["image_grid_thw"] for v in vision]) if vision else None |
| return Batch(ids, mask, torch.tensor(rows), torch.tensor(poss), torch.tensor(codes), torch.tensor(owners), sizes, ros, |
| pixels, grids) |
|
|
|
|
| class DecisionModel(nn.Module): |
| def __init__(self, lm: nn.Module, code_token_ids: Sequence[int], codes: Sequence[str], *, |
| attention: str = "two_way", readout: str = "lm_head", attn_impl: str = "sdpa"): |
| super().__init__() |
| if attention not in ("two_way", "causal"): |
| raise ValueError("attention must be two_way or causal") |
| if attn_impl not in ("sdpa", "varlen"): |
| raise ValueError("attn_impl must be sdpa or varlen") |
| if readout not in ("lm_head", "head"): |
| raise ValueError("readout must be lm_head or head") |
| self.lm = lm |
| self.attention = attention |
| self.readout = readout |
| self.codes = list(codes) |
| self.code_index = {c: i for i, c in enumerate(self.codes)} |
| self.register_buffer("code_token_ids", torch.tensor(list(code_token_ids), dtype=torch.long), persistent=False) |
| rows = self.lm_head_weight()[self.code_token_ids].detach().clone() |
| self.head = nn.Parameter(rows) if readout == "head" else None |
| self.image_encoder: Any = None |
| set_full_attention_causal(lm, attention == "causal") |
| self.attn_impl = attn_impl |
| if attn_impl == "varlen": |
| from glance_runtime import attention as att |
|
|
| att.use(_text_config(lm)) |
|
|
| def lm_head_weight(self) -> torch.Tensor: |
| return self.lm.get_output_embeddings().weight |
|
|
| def code_rows(self) -> torch.Tensor: |
| return self.head if self.head is not None else self.lm_head_weight()[self.code_token_ids] |
|
|
| def hidden(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, pixel_values: torch.Tensor | None = None, |
| image_grid_thw: torch.Tensor | None = None) -> torch.Tensor: |
| core = _core(self.lm) |
| embeds = core.get_input_embeddings()(input_ids) |
| position_ids = None |
| if pixel_values is not None: |
| mm = _multimodal(self.lm) |
| if mm is None: |
| raise ValueError("pixel_values given to a model without a vision encoder") |
| feats = mm.get_image_features(pixel_values, image_grid_thw, return_dict=True).pooler_output |
| feats = torch.cat(list(feats), 0).to(embeds.device, embeds.dtype) |
| image_mask, _ = mm.get_placeholder_mask(input_ids, inputs_embeds=embeds, image_features=feats) |
| embeds = embeds.masked_scatter(image_mask, feats) |
| types = (input_ids == mm.config.image_token_id).int() |
| position_ids, _ = mm.get_rope_index(input_ids, mm_token_type_ids=types, image_grid_thw=image_grid_thw, |
| attention_mask=attention_mask) |
| mask: Any = attention_mask |
| if self.attn_impl == "varlen": |
| mask = varlen_masks(self.lm, embeds, attention_mask, causal=self.attention == "causal") |
| elif self.attention == "two_way": |
| mask = two_way_masks(self.lm, embeds, attention_mask) |
| out = core(inputs_embeds=embeds, attention_mask=mask, position_ids=position_ids, use_cache=False) |
| return out.last_hidden_state |
|
|
| def forward(self, batch: Batch) -> list[torch.Tensor]: |
| """One logit vector per readout, over that question's options in display order.""" |
| h = self.hidden(batch.input_ids, batch.attention_mask, batch.pixel_values, batch.image_grid_thw) |
| hr = h[batch.readout_row, batch.readout_pos] |
| w = self.code_rows()[batch.option_code] |
| flat = (hr[batch.option_readout] * w.to(hr.dtype)).sum(-1) |
| return list(torch.split(flat, batch.sizes)) |
|
|
