modeling_pointsv15_chat.py

from typing import List, Optional, Tuple

import numpy as np
import torch
from PIL import Image
from transformers import GenerationMixin, PreTrainedModel, PreTrainedTokenizer

try:
    from transformers.models.qwen2_vl.image_processing_qwen2_vl import (  # noqa
        Qwen2VLImageProcessor,
    )
    from transformers.models.qwen2_vl.modeling_qwen2_vl import PatchMerger
except ImportError:
    print('Please upgrade transformers to version 4.46.3 or higher')

from .configuration_pointsv15_chat import POINTSV15ChatConfig
from .modeling_llama import CustomLlamaForCausalLM

try:
    from wepoints.models import Qwen2VisionTransformerForNavitPOINTS
except ImportError:
    print('Please install WePOINTS, and refer to https://github.com/WePOINTS/WePOINTS')


class POINTSV15ChatModel(PreTrainedModel, GenerationMixin):
    config_class = POINTSV15ChatConfig
    _no_split_modules = ["CustomLlamaLayer",
                         "Qwen2VisionTransformerPretrainedModel"]

    """Chat model for POINTSv1.5.
    
    Args:
        config (POINTSChatConfigV15): The model config.
    """

    def __init__(self, config: POINTSV15ChatConfig) -> None:
        super().__init__(config)
        self.llm = CustomLlamaForCausalLM(config.llm_config)
        self.vision_encoder = Qwen2VisionTransformerForNavitPOINTS._from_config(  # noqa
            config.vision_config, attn_implementation="flash_attention_2"
        )
        self.vision_projector = PatchMerger(config.llm_config.hidden_size,
                                            context_dim=1280)

    def process_images(self, images: torch.Tensor, 
                       image_grid_thws: List[list]) -> torch.Tensor:
        """Obtain image features from the vision encoder.
        
        Args:
            images (torch.Tensor): The input images.
            image_grid_thws (List[list]): The grid thresholds for the images.

        Returns:
            torch.Tensor: The image features.
        """
        image_features = self.vision_encoder(images, grid_thw=image_grid_thws)
        image_features = self.vision_projector(image_features)
        return image_features
    
    def construct_prompt(self, messages: List[dict],
                         image_processor: Qwen2VLImageProcessor) -> Tuple[str, List[Image.Image], List[list]]: # noqa
        """Construct the prompt for the chat model.

        Args:
            messages (List[dict]): The input messages.
        
        Returns:
            Tuple[str, List[Image.Image], List[list]]: 
                The prompt, images, and image grid shape.
        """
        images = []
        image_grid_thws = []
        reconstructed_messages = []
        for message in messages:
            role = message['role']
            content_from_role = ''
            for item in message['content']:
                if item['type'] == 'text':
                    content_from_role += item['text']
                elif item['type'] == 'image':
                    image_path = item['image']
                    image = Image.open(image_path).convert('RGB')
                    image_data = image_processor(images=image)
                    pixel_values = image_data['pixel_values']
                    image_grid_thw = image_data['image_grid_thw']
                    images.extend(pixel_values)
                    image_grid_thws.append(image_grid_thw)
                    seq_len = int(image_grid_thw[0][1] * image_grid_thw[0][2] / 4) # noqa
                    content_from_role += '<|vision_start|>' + '<|image_pad|>' * seq_len + '<|vision_end|>' + '\n' # noqa
            reconstructed_messages.append({
                'role': role,
                'content': content_from_role
            })
        prompt = self.apply_chat_template(reconstructed_messages)
        return prompt, images, image_grid_thws
    
    def apply_chat_template(self, messages: List[dict]) -> str:
        """Apply the chat template to the input messages.

        Args:
            messages (List[dict]): The input messages.
        
        Returns:
            str: The prompt.
        """
        role_prefix_mapping = {
            'user': '<|im_start|>user\n',
            'assistant': '<|im_start|>assistant\n'
        }
        role = 'user'
        prompt = ''
        for message in messages:
            role = message['role']
            content = message['content']
            prompt += role_prefix_mapping[role] + content + '<|im_end|>\n'
        if role == 'user':
            prompt += '<|im_start|>assistant\n'
        return prompt

    @torch.no_grad()
    def chat(self, 
             messages: List[dict],
             tokenizer: PreTrainedTokenizer,
             image_processor: object,
             generation_config: dict = None) -> str:
        """Generate a response to the input prompt.

        Args:
            messages (List[dict]): The input messages.
            tokenizer (PreTrainedTokenizer): The tokenizer to use.
            image_processor (object): The image processor to use.
            generation_config (dict, optional): The generation config. 
                Defaults to None.
        Returns:
            str: The generated response.
        """
        prompt, images, image_grid_thws = self.construct_prompt(
            messages, image_processor
        )
        images = np.array(images)
        images = torch.from_numpy(images).to(self.vision_encoder.device).to(self.vision_encoder.dtype) # noqa
        image_grid_thws = np.concatenate(image_grid_thws, axis=0)
        image_grid_thws = (
            torch.from_numpy(image_grid_thws)
            .cuda()
            .long()
        )
        image_features = self.vision_encoder(images, grid_thw=image_grid_thws)
        image_features = self.vision_projector(image_features)
        model_inputs = tokenizer(prompt, return_tensors='pt')
        input_ids = model_inputs['input_ids'].to(self.device)
        attention_mask = model_inputs['attention_mask'].to(self.device)
        # stop token
        eos_token_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
        # image token
        image_token_id = tokenizer.convert_tokens_to_ids("<|image_pad|>")
        generation_config.update(
            {
                'eos_token_id': eos_token_id,
            }
        )
        outputs = self.generate(
            input_ids=input_ids,
            image_grid_thws=image_grid_thws,
            attention_mask=attention_mask,
            image_features=[image_features],
            image_token_id=image_token_id,
            **generation_config
        )
        response = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
        return response
      
    def _split_input_ids(self, input_ids, special_token):
        special_pos = input_ids == special_token
        pos = (special_pos[:-1] != special_pos[1:]).nonzero() + 1
        if pos.shape[0] % 2 != 0:
            pos = torch.cat([torch.tensor([[0]]).to(pos.device), pos])
        pos = pos.reshape(-1, 2).tolist()
        return pos

    def generate(self,
                 input_ids: torch.LongTensor,
                 image_grid_thws: torch.LongTensor,
                 attention_mask: torch.LongTensor,
                 image_features: List[torch.Tensor],
                 image_token_id: int,
                 generation_config: Optional[dict] = None,
                 output_hidden_states: Optional[bool] = None,
                 return_dict: Optional[bool] = None,
                 **generate_kwargs) -> torch.LongTensor:
        input_embeddings = self.llm.lm.embed_in(input_ids)
        batch_size = input_ids.shape[0]
        assert len(image_features) == batch_size
        for i in range(batch_size):
            pos = self._split_input_ids(input_ids[i], image_token_id)
            assert len(pos) == len(image_grid_thws)
            image_pos = [
                int(image_grid_thw[1] * image_grid_thw[2] / 4)
                for image_grid_thw in image_grid_thws
            ]
            image_pos.insert(0, 0)
            image_pos = np.cumsum(image_pos)
            for j, (start, end) in enumerate(pos):
                input_embeddings[i, start:end] = \
                    image_features[i][image_pos[j]:image_pos[j+1]]
        outputs = self.llm.generate(
            inputs_embeds=input_embeddings,
            attention_mask=attention_mask,
            generation_config=generation_config,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            use_cache=True,
            **generate_kwargs
        )
        return outputs