modeling_vectorllm.py

from typing import List, Optional, Tuple, Union

import torch
from torch import nn
from torch.nn import CrossEntropyLoss
from transformers import GenerationConfig, Qwen3Config, Qwen3ForCausalLM
from transformers.activations import ACT2FN
from transformers.modeling_outputs import CausalLMOutputWithPast
from transformers.modeling_utils import PreTrainedModel

from .adaptor_base import *  # noqa: F401,F403
from .adaptor_generic import *  # noqa: F401,F403
from .adaptor_mlp import *  # noqa: F401,F403
from .adaptor_registry import *  # noqa: F401,F403
from .cls_token import *  # noqa: F401,F403
from .configuration_vectorllm import ProjectorConfig, VectorLLMConfig
from .common import *  # noqa: F401,F403
from .dinov2_arch import *  # noqa: F401,F403
from .dual_hybrid_vit import *  # noqa: F401,F403
from .enable_cpe_support import *  # noqa: F401,F403
from .enable_spectral_reparam import *  # noqa: F401,F403
from .eradio_model import *  # noqa: F401,F403
from .extra_models import *  # noqa: F401,F403
from .extra_timm_models import *  # noqa: F401,F403
from .feature_normalizer import *  # noqa: F401,F403
from .forward_intermediates import *  # noqa: F401,F403
from .hf_model import RADIOConfig as HFRADIOConfig, RADIOModel as HFRADIOModel
from .input_conditioner import *  # noqa: F401,F403
from .open_clip_adaptor import *  # noqa: F401,F403
from .radio_model import *  # noqa: F401,F403
from .vit_patch_generator import *  # noqa: F401,F403
from .vitdet import *  # noqa: F401,F403


IGNORE_INDEX = -100


def prepare_inputs_labels_for_multimodal_vectorllm(
    llm,
    input_ids: torch.LongTensor = None,
    position_ids: Optional[torch.LongTensor] = None,
    attention_mask: Optional[torch.Tensor] = None,
    past_key_values: Optional[List[torch.FloatTensor]] = None,
    labels: Optional[torch.LongTensor] = None,
    pixel_values: Optional[torch.FloatTensor] = None,
    pixel_token_id=None,
):
    if pixel_values is None:
        return {
            "input_ids": input_ids,
            "position_ids": position_ids,
            "attention_mask": attention_mask,
            "past_key_values": past_key_values,
            "inputs_embeds": None,
            "labels": labels,
        }

    original_labels = labels
    original_position_ids = position_ids
    original_attention_mask = attention_mask

    if attention_mask is None:
        attention_mask = torch.ones_like(input_ids, dtype=torch.bool)
    else:
        attention_mask = attention_mask.bool()
    if position_ids is None:
        position_ids = torch.arange(
            0, input_ids.shape[1], dtype=torch.long, device=input_ids.device
        ).unsqueeze(0).expand(input_ids.shape[0], -1)
    if labels is None:
        labels = torch.full_like(input_ids, IGNORE_INDEX)

    inputs_embeds = llm.get_input_embeddings()(input_ids)
    inputs_embeds = inputs_embeds.clone()
    labels = labels.clone()

    if pixel_values.ndim != 3:
        raise ValueError(f"Expected pixel_values to have shape [B, N, C], got {tuple(pixel_values.shape)}")

    for batch_idx in range(input_ids.shape[0]):
        replace_positions = torch.where(input_ids[batch_idx] == pixel_token_id)[0]
        if replace_positions.numel() == 0:
            continue
        if replace_positions.numel() != pixel_values.shape[1]:
            raise ValueError(
                "The number of image placeholder tokens does not match the projected visual tokens: "
                f"{replace_positions.numel()} vs {pixel_values.shape[1]}"
            )
        inputs_embeds[batch_idx, replace_positions] = pixel_values[batch_idx].to(inputs_embeds.dtype)
        labels[batch_idx, replace_positions] = IGNORE_INDEX

    return {
        "input_ids": None,
        "position_ids": None if original_position_ids is None else position_ids,
        "attention_mask": None if original_attention_mask is None else attention_mask.to(dtype=original_attention_mask.dtype),
        "past_key_values": past_key_values,
        "inputs_embeds": inputs_embeds,
        "labels": None if original_labels is None else labels,
    }


class ProjectorModel(PreTrainedModel):
    config_class = ProjectorConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True

    def __init__(self, config: ProjectorConfig) -> None:
        super().__init__(config)
        self.gradient_checkpointing = False
        modules = [
            nn.Linear(config.visual_hidden_size, config.llm_hidden_size, bias=config.bias)
        ]
        for _ in range(1, config.depth):
            modules.append(ACT2FN[config.hidden_act])
            modules.append(
                nn.Linear(config.llm_hidden_size, config.llm_hidden_size, bias=config.bias)
            )
        self.model = nn.Sequential(*modules)

    def forward(self, x):
        if self.gradient_checkpointing and self.training:
            return torch.utils.checkpoint.checkpoint(self.model, x)
        return self.model(x)


class VectorLLMForCausalLM(PreTrainedModel):
    config_class = VectorLLMConfig
    main_input_name = "pixel_values"
    base_model_prefix = "model"
    supports_gradient_checkpointing = True

    def __init__(
        self,
        config: VectorLLMConfig,
        vision_model=None,
        language_model=None,
        projector=None,
        pos_embeds=None,
    ):
        super().__init__(config)

        if vision_model is not None:
            self.vision_model = vision_model
        else:
            self.vision_model = HFRADIOModel(HFRADIOConfig(**config.vision_config))
            target_dtype = getattr(torch, config.vision_torch_dtype, None)
            if target_dtype is not None:
                self.vision_model = self.vision_model.to(dtype=target_dtype)

        if language_model is not None:
            self.language_model = language_model
        else:
            self.language_model = Qwen3ForCausalLM(Qwen3Config(**config.llm_config))

        if projector is not None:
            self.projector = projector
        else:
            self.projector = ProjectorModel(ProjectorConfig(**config.projector_config))

        width = config.regression_size[0] // config.patch_size
        height = config.regression_size[1] // config.patch_size
        n_pos = width * height
        if pos_embeds is not None:
            self.visual_pos_embeddings = pos_embeds
        else:
            self.visual_pos_embeddings = nn.Embedding(n_pos, config.vision_hidden_size)

        self.pixel_idx = config.pixel_idx
        self.num_cls_register_tokens = config.num_cls_register_tokens

    @property
    def lm_head(self):
        return self.language_model.get_output_embeddings()

    def get_input_embeddings(self):
        return self.language_model.get_input_embeddings()

    def get_output_embeddings(self):
        return self.language_model.get_output_embeddings()

    def extract_feature(self, pixel_values):
        summary, features = self.vision_model(pixel_values.to(self.vision_model.dtype))
        del summary
        pos_embed = self.visual_pos_embeddings.weight.unsqueeze(0)
        pos_embed = pos_embed.repeat(features.shape[0], 1, 1)
        features = features + pos_embed
        features = features.to(self.projector.dtype)
        return self.projector(features)

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values=None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        cache_position: Optional[torch.LongTensor] = None,
        pixel_values: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
    ):
        if inputs_embeds is None and pixel_values is not None:
            if isinstance(pixel_values, list):
                pixel_values = [item.unsqueeze(0) if item.ndim == 3 else item for item in pixel_values]
                pixel_values = torch.cat(pixel_values, dim=0)
            pixel_values = pixel_values.to(self.device)
            projected = self.extract_feature(pixel_values)
            llm_inputs = prepare_inputs_labels_for_multimodal_vectorllm(
                llm=self.language_model,
                input_ids=input_ids,
                position_ids=position_ids,
                attention_mask=attention_mask,
                past_key_values=past_key_values,
                labels=labels,
                pixel_values=projected,
                pixel_token_id=self.pixel_idx,
            )
            inputs_embeds = llm_inputs["inputs_embeds"]
            attention_mask = llm_inputs["attention_mask"]
            position_ids = llm_inputs["position_ids"]
            labels = llm_inputs["labels"]
            input_ids = llm_inputs["input_ids"]

        outputs = self.language_model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        logits = outputs.logits
        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss_fct = CrossEntropyLoss()
            shift_logits = shift_logits.view(-1, self.language_model.config.vocab_size)
            shift_labels = shift_labels.view(-1).to(shift_logits.device)
            loss = loss_fct(shift_logits, shift_labels)

        if not return_dict:
            output = (logits,) + outputs[1:]
            return (loss,) + output if loss is not None else output

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

    @torch.no_grad()
    def generate(
        self,
        pixel_values: Optional[torch.FloatTensor] = None,
        input_ids: Optional[torch.FloatTensor] = None,
        attention_mask: Optional[torch.LongTensor] = None,
        generation_config: Optional[GenerationConfig] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict_in_generate: Optional[bool] = None,
        **generate_kwargs,
    ) -> torch.LongTensor:
        if pixel_values is not None:
            if isinstance(pixel_values, list):
                pixel_values = [item.unsqueeze(0) if item.ndim == 3 else item for item in pixel_values]
                pixel_values = torch.cat(pixel_values, dim=0)
            pixel_values = pixel_values.to(self.device)
            input_ids = input_ids.to(self.device)
            input_embeds = self.language_model.get_input_embeddings()(input_ids)
            projected = self.extract_feature(pixel_values).to(input_embeds.dtype)
            batch, seqlen, channels = input_embeds.shape
            flat_embeds = input_embeds.reshape(batch * seqlen, channels)
            selected = input_ids.reshape(batch * seqlen) == self.pixel_idx
            flat_embeds[selected] = projected.reshape(-1, channels).to(flat_embeds.device)
            input_embeds = flat_embeds.reshape(batch, seqlen, channels)
        else:
            input_embeds = self.language_model.get_input_embeddings()(input_ids.to(self.device))

        outputs = self.language_model.generate(
            inputs_embeds=input_embeds,
            attention_mask=attention_mask.to(self.device) if attention_mask is not None else None,
            generation_config=generation_config,
            output_hidden_states=output_hidden_states,
            return_dict_in_generate=return_dict_in_generate,
            **generate_kwargs,
        )
        return outputs