vllm_support/src/kalm_t5gemma2_vllm_plugin/modeling.py

from __future__ import annotations

import math
import os
from typing import Iterable

import torch
import torch.nn as nn
from transformers.modeling_outputs import BaseModelOutput
from vllm.model_executor.layers.pooler import DispatchPooler
from vllm.multimodal import MULTIMODAL_REGISTRY

from .constants import (
    DEFAULT_DECODER_PAD_TO_MULTIPLE_OF,
    DEFAULT_ENCODER_CHUNK_SIZE,
    NO_TOKEN_ID,
    TEXT_MODALITY,
    YES_TOKEN_ID,
)
from .modeling_score import T5Gemma2ForScoreClassification
from .processing import (
    TextEncoderDummyInputsBuilder,
    TextEncoderProcessingInfo,
    TextEncoderProcessor,
)


def _as_token_rows(value: object) -> list[torch.Tensor]:
    if isinstance(value, torch.Tensor):
        if value.ndim == 1:
            return [value]
        if value.ndim == 2:
            return [row for row in value]
        raise ValueError(f"encoder_input_ids must be 1D/2D, got {value.shape}.")
    if isinstance(value, list):
        return [
            item.flatten()
            if isinstance(item, torch.Tensor)
            else torch.tensor(item, dtype=torch.long)
            for item in value
        ]
    raise TypeError(f"Unsupported encoder_input_ids type: {type(value)!r}")


def _split_by_position_zero(
    input_ids: torch.Tensor,
    positions: torch.Tensor,
) -> tuple[list[torch.Tensor], list[int]]:
    flat_ids = input_ids.flatten()
    starts = (positions.flatten() == 0).nonzero(as_tuple=False).flatten().tolist()
    if not starts:
        return [], []
    starts.append(int(flat_ids.numel()))
    rows: list[torch.Tensor] = []
    last_indices: list[int] = []
    for start, end in zip(starts[:-1], starts[1:]):
        if end > start:
            rows.append(flat_ids[start:end])
            last_indices.append(end - 1)
    return rows, last_indices


def _debug(message: str) -> None:
    if os.environ.get("KALM_VLLM_DEBUG") == "1":
        print(f"[kalm-vllm-debug] {message}", flush=True)


@MULTIMODAL_REGISTRY.register_processor(
    TextEncoderProcessor,
    info=TextEncoderProcessingInfo,
    dummy_inputs=TextEncoderDummyInputsBuilder,
)
class T5Gemma2VllmScoreClassification(nn.Module):
    is_pooling_model = True
    supports_multimodal = True
    score_type = "cross-encoder"
    attn_type = "encoder_decoder"
    default_seq_pooling_type = "LAST"
    default_tok_pooling_type = "ALL"

    def __init__(self, *, vllm_config, prefix: str = "") -> None:
        super().__init__()
        self.vllm_config = vllm_config
        self.model_config = vllm_config.model_config
        self.config = self.model_config.hf_config
        self.config.num_labels = 1
        self.config.yes_token_id = int(
            getattr(self.config, "yes_token_id", YES_TOKEN_ID)
        )
        self.config.no_token_id = int(
            getattr(self.config, "no_token_id", NO_TOKEN_ID)
        )
        self.config.encoder_chunk_size = getattr(
            self.config,
            "encoder_chunk_size",
            DEFAULT_ENCODER_CHUNK_SIZE,
        )
        self.encoder_chunk_size = self.config.encoder_chunk_size
        self.config.decoder_pad_to_multiple_of = int(
            getattr(
                self.config,
                "decoder_pad_to_multiple_of",
                DEFAULT_DECODER_PAD_TO_MULTIPLE_OF,
            )
        )
        self.decoder_pad_to_multiple_of = self.config.decoder_pad_to_multiple_of
        self.pad_token_id = int(getattr(self.config, "pad_token_id", 0) or 0)

        self.score_model = T5Gemma2ForScoreClassification.from_pretrained(
            self.model_config.model,
            trust_remote_code=True,
            dtype=self.model_config.dtype,
        )
        self.score_model.config.yes_token_id = self.config.yes_token_id
        self.score_model.config.no_token_id = self.config.no_token_id
        self.score_model.config.num_labels = 1
        self.score_model.config.encoder_chunk_size = self.encoder_chunk_size
        self.score_model.yes_token_id = self.config.yes_token_id
        self.score_model.no_token_id = self.config.no_token_id
        self.score_model.encoder_chunk_size = self.encoder_chunk_size
        self.score_model._validate_score_config()
        self.score_model.eval()

        pooler_config = self.model_config.pooler_config
        assert pooler_config is not None
        self.pooler = DispatchPooler.for_seq_cls(pooler_config)

    def get_language_model(self):
        return self

    def get_num_mm_encoder_tokens(self, num_tokens: int) -> int:
        if self.encoder_chunk_size is None:
            return int(num_tokens)
        return int(math.ceil(num_tokens / int(self.encoder_chunk_size)))

    def embed_input_ids(self, input_ids: torch.Tensor, *args, **kwargs) -> torch.Tensor:
        return self.score_model.get_input_embeddings()(input_ids)

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        # The semantic wrapper loads the same checkpoint directly.
        for _ in weights:
            pass
        return set(self.state_dict().keys())

    def embed_multimodal(self, **kwargs: object) -> list[torch.Tensor]:
        if "encoder_input_ids" not in kwargs:
            raise ValueError(f"Missing {TEXT_MODALITY} encoder_input_ids.")
        rows = _as_token_rows(kwargs["encoder_input_ids"])
        device = next(self.score_model.parameters()).device
        outputs: list[torch.Tensor] = []
        with torch.inference_mode():
            for row in rows:
                input_ids = row.to(device=device, dtype=torch.long).unsqueeze(0)
                attention_mask = torch.ones_like(input_ids)
                raw_encoder_outputs = self.score_model.get_encoder()(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    return_dict=True,
                )
                hidden = raw_encoder_outputs.last_hidden_state
                if self.encoder_chunk_size is not None:
                    hidden, _ = self.score_model._pool_encoder_chunks(
                        hidden,
                        attention_mask,
                        int(self.encoder_chunk_size),
                    )
                item = hidden.squeeze(0).contiguous()
                _debug(f"encoder output shape={tuple(item.shape)}")
                outputs.append(item)
        return outputs

    def _pad_decoder_rows(
        self,
        rows: list[torch.Tensor],
        device: torch.device,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        max_len = max(int(row.numel()) for row in rows)
        if self.decoder_pad_to_multiple_of > 1:
            multiple = self.decoder_pad_to_multiple_of
            max_len = int(math.ceil(max_len / multiple) * multiple)
        batch = torch.full(
            (len(rows), max_len),
            self.pad_token_id,
            dtype=torch.long,
            device=device,
        )
        mask = torch.zeros_like(batch)
        for index, row in enumerate(rows):
            row = row.to(device=device, dtype=torch.long)
            length = int(row.numel())
            batch[index, :length] = row
            mask[index, :length] = 1
        return batch, mask

    @staticmethod
    def _pad_encoder_outputs(
        encoder_outputs: list[torch.Tensor],
        device: torch.device,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        max_len = max(int(item.shape[0]) for item in encoder_outputs)
        hidden_size = int(encoder_outputs[0].shape[-1])
        batch = torch.zeros(
            (len(encoder_outputs), max_len, hidden_size),
            dtype=encoder_outputs[0].dtype,
            device=device,
        )
        mask = torch.zeros(
            (len(encoder_outputs), max_len),
            dtype=torch.long,
            device=device,
        )
        for index, item in enumerate(encoder_outputs):
            item = item.to(device=device)
            length = int(item.shape[0])
            batch[index, :length] = item
            mask[index, :length] = 1
        return batch, mask

    def forward(
        self,
        input_ids: torch.Tensor | None,
        positions: torch.Tensor,
        intermediate_tensors=None,
        inputs_embeds: torch.Tensor | None = None,
        encoder_outputs: list[torch.Tensor] | torch.Tensor | None = None,
        **kwargs,
    ) -> torch.Tensor:
        if input_ids is None:
            raise ValueError("Decoder input_ids are required.")
        decoder_rows, last_indices = _split_by_position_zero(input_ids, positions)
        hidden = input_ids.new_zeros((input_ids.numel(), 1), dtype=torch.float32)
        if not decoder_rows:
            return hidden
        if encoder_outputs is None:
            return hidden

        encoder_list = (
            [item for item in encoder_outputs]
            if isinstance(encoder_outputs, torch.Tensor)
            else list(encoder_outputs)
        )
        if len(encoder_list) != len(decoder_rows):
            raise ValueError(
                "Mismatched encoder/decoder batch sizes: "
                f"{len(encoder_list)} vs {len(decoder_rows)}."
            )

        device = next(self.score_model.parameters()).device
        decoder_batch, decoder_mask = self._pad_decoder_rows(decoder_rows, device)
        encoder_batch, encoder_mask = self._pad_encoder_outputs(encoder_list, device)
        with torch.inference_mode():
            outputs = self.score_model(
                encoder_outputs=BaseModelOutput(last_hidden_state=encoder_batch),
                attention_mask=encoder_mask,
                decoder_input_ids=decoder_batch,
                decoder_attention_mask=decoder_mask,
            )
            margins = outputs.logits.squeeze(-1).to(hidden.device, dtype=hidden.dtype)
            _debug(f"margins={margins.float().tolist()}")
        for row_index, last_index in enumerate(last_indices):
            hidden[last_index, 0] = margins[row_index]
        return hidden


__all__ = ["T5Gemma2VllmScoreClassification"]