| from typing import Optional |
| import torch |
| from torch import nn |
| from transformers import PreTrainedModel |
| from transformers.models.qwen3_vl import Qwen3VLModel |
| from transformers.utils import logging |
|
|
| from .configuration_ops_colqwen3 import OpsColQwen3Config |
|
|
| logger = logging.get_logger(__name__) |
|
|
|
|
| class OpsColQwen3PreTrainedModel(PreTrainedModel): |
| config_class = OpsColQwen3Config |
| base_model_prefix = "ops_colqwen3" |
| supports_gradient_checkpointing = True |
| _no_split_modules = ["Qwen3VLVisionBlock", "Qwen3DecoderLayer"] |
| _skip_keys_device_placement = "past_key_values" |
| _supports_flash_attn_2 = True |
| _supports_sdpa = True |
| _supports_cache_class = True |
|
|
|
|
| class OpsColQwen3Model(OpsColQwen3PreTrainedModel): |
| _checkpoint_conversion_mapping = { |
| r"^language_model": r"qwen3vl.language_model", |
| r"^visual": "qwen3vl.visual", |
| } |
|
|
| def __init__(self, config: OpsColQwen3Config): |
| super().__init__(config) |
| self.config = config |
|
|
| self.qwen3vl = Qwen3VLModel(config) |
| self.dims = config.text_config.hidden_size |
| self.custom_text_proj = nn.Linear(config.text_config.hidden_size, self.dims) |
|
|
| self.mask_non_image_embeddings = config.mask_non_image_embeddings |
| self.post_init() |
|
|
| @classmethod |
| def from_pretrained(cls, *args, config: Optional[OpsColQwen3Config] = None, **kwargs): |
| key_mapping = kwargs.pop("key_mapping", None) |
| if key_mapping is None: |
| key_mapping = getattr(cls, "_checkpoint_conversion_mapping", None) |
| dims = None |
| if 'dims' in kwargs: |
| dims = kwargs.pop('dims') |
| elif config is not None: |
| dims = config.dims |
|
|
| model = super().from_pretrained(*args, config=config, **kwargs, key_mapping=key_mapping) |
| if dims is not None: |
| model.dims = dims |
| return model |
|
|
| def forward(self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, pixel_values: Optional[torch.Tensor] = None, image_grid_thw: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor: |
| has_pixel_values = pixel_values is not None |
|
|
| if has_pixel_values: |
| if image_grid_thw is None: |
| raise ValueError("`image_grid_thw` must be provided when `pixel_values` is passed.") |
| if not torch.is_tensor(image_grid_thw): |
| image_grid_thw = torch.as_tensor(image_grid_thw, device=pixel_values.device) |
|
|
| offsets = image_grid_thw.prod(dim=1) |
| unpadded = [pixel_sequence[: int(offset.item())] for pixel_sequence, offset in zip(pixel_values, offsets)] |
| pixel_values = torch.cat(unpadded, dim=0) if unpadded else None |
|
|
| outputs = self.qwen3vl( |
| input_ids=input_ids, |
| attention_mask=attention_mask, |
| pixel_values=pixel_values, |
| image_grid_thw=image_grid_thw, |
| use_cache=False, |
| output_hidden_states=True, |
| return_dict=True, |
| ) |
|
|
| last_hidden_states = outputs.last_hidden_state |
| proj = self.custom_text_proj(last_hidden_states) |
|
|
| if self.dims < self.config.text_config.hidden_size: |
| proj = proj[..., : self.dims] |
|
|
| proj = proj / proj.norm(dim=-1, keepdim=True) |
|
|
| if attention_mask is not None: |
| proj = proj * attention_mask.unsqueeze(-1) |
|
|
| if has_pixel_values and self.mask_non_image_embeddings and input_ids is not None: |
| image_mask = (input_ids == self.config.image_token_id).unsqueeze(-1) |
| proj = proj * image_mask |
|
|
| return proj |
|
|
| @property |
| def patch_size(self) -> int: |
| return self.qwen3vl.visual.config.patch_size |
|
|
| @property |
| def spatial_merge_size(self) -> int: |
| return self.qwen3vl.visual.config.spatial_merge_size |
|
|
