processing_mplugowl3.py

# coding=utf-8
# Copyright 2024 The HuggingFace Inc. team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Processor class for mPLUGOwl3.
"""

from typing import List, Optional, Union, Dict, Any
import warnings
import torch
import re

from transformers.image_processing_utils import BatchFeature
from transformers.image_utils import ImageInput
from transformers.processing_utils import ProcessorMixin
from transformers.tokenization_utils_base import PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
from transformers.utils import TensorType, requires_backends, is_torch_dtype, is_torch_device

from .image_processing_mplugowl3 import mPLUGOwl3BatchFeature, mPLUGOwl3ImageProcessor

OWL_MEDIA_TOKEN=['<|image|>']

class MediaIndicesHelper():
 def __init__(self, tokenizer) -> None:
 self.media_position = []
 self.tokenizer = tokenizer
 
 
 def has_media(self, text, media_tokens=None):
 if media_tokens is None:
 media_tokens = OWL_MEDIA_TOKEN
 has_media_flag = any([media_token == text for media_token in media_tokens])
 if any([media_token in text for media_token in media_tokens]):
 # 不允许出现text中包含media token但是不仅仅是media token。 media token必须单独为一个chunk 
 assert has_media_flag, text
 return has_media_flag
 
 def add_media(self, text_chunk, text=None, tokenize_fn=None):
 
 # cross
 assert tokenize_fn is not None
 assert text is not None
 assert text in OWL_MEDIA_TOKEN
 media_token_ids = tokenize_fn(text)
 start = len(text_chunk)
 end = start + len(media_token_ids)
 self.media_position.append([start, end])
 text_chunk.extend(media_token_ids)
 return len(media_token_ids)

 def cal_media_offset(self, input_ids):
 if len(self.media_position) == 0:
 return torch.ones_like(input_ids)*(-1000000)

 media_starts = torch.tensor([_[0] for _ in self.media_position]).reshape(1,-1)
 rng = torch.arange(input_ids.shape[0]).reshape(-1,1)
 matrix = (rng > media_starts).sum(dim=1)
 
 return matrix
 
 def len_images(self,):
 return len(self.media_position)

class mPLUGOwl3Processor(ProcessorMixin):
 r"""
 Args:
 image_processor ([`mPLUGOwl3ImageProcessor`], *optional*):
 The image processor is a required input.
 tokenizer ([`LlamaTokenizerWrapper`], *optional*):
 The tokenizer is a required input.
 """
 attributes = ["image_processor", "tokenizer"]
 image_processor_class = "AutoImageProcessor"
 tokenizer_class = "AutoTokenizer"

 def __init__(self, image_processor: mPLUGOwl3ImageProcessor = None, tokenizer=None, prompt_style='chatml', inference_mode=True, addition_eod="<|endoftext|>"):
 super().__init__(image_processor, tokenizer)
 self.image_processor: mPLUGOwl3ImageProcessor
 self.prompt_style = prompt_style
 self.inference_mode = inference_mode
 self.media_tokens = ["<|image|>"]
 self.addition_eod = addition_eod

 def build_text_qwen(self, messages):
 # role should be within ['system', 'user', 'assistant']
 im_start, im_end = '<|im_start|>', '<|im_end|>'
 
 text = []
 for num_turn, message in enumerate(messages):
 if num_turn == 0 and message['role'] != 'system':
 if self.prompt_style != 'plain':
 text.append({
 "text": f"{im_start}system\n{im_end}",
 "label": 0
 })
 if message['role'] == 'system':
 if self.prompt_style != 'plain':
 text.append({
 "text": f"{im_start}system\n{message['content']}{im_end}",
 "label": 0
 })
 elif message['role'] == 'user':
 if self.prompt_style != 'plain':
 content = f"\n{im_start}user\n{message['content']}{im_end}"
 else:
 content = message['content']
 pattern = '|'.join(map(re.escape, self.media_tokens))
 chunk_strs = re.split(f'({pattern})', content)
 for chunk_str in chunk_strs:
 text.append({
 "text": chunk_str,
 "label": 0
 })
 
 elif message['role'] == 'assistant':
 if self.prompt_style != 'plain':
 text.append({"text": f"\n{im_start}assistant\n", "label": 0})
 text.append({"text": f"{message['content']}{im_end}", "label": 1})
 else:
 text.append({"text": f"{message['content']}", "label": 1})
 text.append({"text": self.addition_eod, "label": 1})
 else:
 raise NotImplementedError
 if self.inference_mode:
 while text and text[-1]['label']==1: # 只要列表非空且最后一个元素满足条件
 text.pop() # 就移除最后一个元素
 return text

 def wrapped_tokenize(self, text):
 return self.tokenizer(text).input_ids

 def encode_text_sft(self, texts):
 # output enc_chunk
 
 enc_chunk = []
 label_chunk = []
 enc_length = 0

 num_images = 0

 media_helper = MediaIndicesHelper(tokenizer=self.tokenizer)
 for current_ti, text_chunk in enumerate(texts):
 
 text = text_chunk["text"]
 label = text_chunk["label"]

 if not media_helper.has_media(text):
 curr_chunk=self.wrapped_tokenize(text)
 if label == 1:
 enc_length += len(curr_chunk)
 enc_chunk += curr_chunk
 label_chunk += [label] * len(curr_chunk)
 else:
 
 enc_length += len(curr_chunk)
 enc_chunk += curr_chunk
 label_chunk += [label] * len(curr_chunk)
 # For media tokens
 else:
 
 add_length = media_helper.add_media(
 enc_chunk, 
 text=text, 
 tokenize_fn=self.wrapped_tokenize)
 enc_length += add_length
 label_chunk += [label] * add_length
 # enc_chunk.extend([self.media_tokens[text]] * self.media_lengths[text])
 # enc_length += self.media_lengths[text]
 # label_chunk += [label] * self.media_lengths[text]
 num_images += 1

 enc_chunk = torch.tensor(enc_chunk).long()
 media_offset = []
 media_before = 0
 for i,_ in enumerate([media_helper]):
 mo = _.cal_media_offset(enc_chunk)
 media_offset.append(torch.cat([(torch.ones(mo.shape[0],1)*media_before).long().to(mo.device), (mo+media_before).unsqueeze(1)], dim=1)) # L 2

 media_before += _.len_images()
 media_offset = torch.stack(media_offset, dim=0)
 return {
 'input_ids': enc_chunk.unsqueeze(0), 
 'media_offset': media_offset,
 }


 def __call__(
 self,
 messages,
 images: ImageInput = None,
 videos = None,
 max_length: Optional[int] = None,
 cut_enable=True,
 return_tensors: Optional[Union[str, TensorType]] = TensorType.PYTORCH,
 **kwargs
 ) -> mPLUGOwl3BatchFeature:
 if videos is not None and len(videos)>0:
 cut_enable=False
 assert images is None or len(images)==0, "We do not support image video interleaved yet"
 video_ptr = 0
 for message in messages:
 text_list = message['content'].split('<|video|>')
 text = text_list[0]
 for next_text in text_list[1:]:
 text += '<|image|>'*len(videos[video_ptr])
 text += next_text
 video_ptr += 1
 message['content'] = text
 images = [frame for video in videos for frame in video ]
 self.check_media(images, messages)
 if images is not None:
 image_inputs = self.image_processor(images, cut_enable=cut_enable, return_tensors=return_tensors)

 if image_inputs.get('cut_shape',None) is not None:
 cut_shape = image_inputs['cut_shape']
 image_token_ptr = 0
 for message in messages:
 text_list = message['content'].split('<|image|>')
 text = text_list[0]
 for next_text in text_list[1:]:
 text += self.image_processor.cut_prompt_template(img_token='<|image|>', h=cut_shape[image_token_ptr][0], w=cut_shape[image_token_ptr][1])
 text += next_text
 image_token_ptr += 1
 if self.image_processor.add_global:
 image_token_ptr += 1 # skip the global image which has been include by cut_prompt_template
 message['content'] = text
 
 
 # text = ''.join([_['text'] for _ in text])
 text = self.build_text_qwen(messages)
 model_inputs = self.encode_text_sft(text)
 
 if images is not None:
 model_inputs.update(image_inputs.data)
 if 'cut_shape' in model_inputs:
 model_inputs.pop('cut_shape')
 if 'cut_shape_indices' in model_inputs:
 model_inputs.pop('cut_shape_indices')
 return mPLUGOwl3BatchFeature(model_inputs)
 
 def check_media(self, images, messages):
 media_num = 0 if images is None else len(images)
 media_count = sum([message['content'].count('<|image|>') for message in messages])
 assert media_num == media_count


 # Copied from transformers.models.clip.processing_clip.CLIPProcessor.batch_decode with CLIP->Llama
 def batch_decode(self, *args, **kwargs):
 """
 This method forwards all its arguments to LlamaTokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please
 refer to the docstring of this method for more information.
 """
 output_ids = args[0]
 result_text = []
 for result in output_ids:
 result = result[result != 0]
 if result[0] == self.tokenizer.bos_id:
 result = result[1:]
 if result[-1] == self.tokenizer.eos_id:
 result = result[:-1]
 result_text.append(self.tokenizer.decode(result, *args[1:], **kwargs).strip())
 return result_text
 # return self.tokenizer.batch_decode(*args, **kwargs)
 
 # Copied from transformers.models.clip.processing_clip.CLIPProcessor.decode with CLIP->Llama
 def decode(self, *args, **kwargs):
 """
 This method forwards all its arguments to LlamaTokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to
 the docstring of this method for more information.
 """
 result = args[0]
 result = result[result != 0]
 if result[0] == self.tokenizer.bos_id:
 result = result[1:]
 if result[-1] == self.tokenizer.eos_id or (hasattr(self.tokenizer, "eot_id") and result[-1] == self.tokenizer.eot_id):
 result = result[:-1]
 return self.tokenizer.decode(result, *args[1:], **kwargs).strip()

 def _convert(
 self, input_str, max_inp_length: Optional[int] = None
 ):
 if self.version > 2.5 or not getattr(self.tokenizer, "add_bos_token", False):
 input_ids = self.tokenizer.encode(input_str)
 else:
 input_ids = [self.tokenizer.bos_id] + self.tokenizer.encode(input_str)
 if max_inp_length is not None:
 input_ids = input_ids[:max_inp_length]
 input_ids = torch.tensor(input_ids, dtype=torch.int32)

 start_cond = (input_ids == self.tokenizer.im_start_id) | (input_ids == self.tokenizer.slice_start_id)
 end_cond = (input_ids == self.tokenizer.im_end_id) | (input_ids == self.tokenizer.slice_end_id)

 image_start_tokens = torch.where(start_cond)[0]
 image_start_tokens += 1
 image_end_tokens = torch.where(end_cond)[0]

 valid_image_nums = max(len(image_start_tokens), len(image_end_tokens))

 image_bounds = torch.hstack(
 [
 image_start_tokens[:valid_image_nums].unsqueeze(-1),
 image_end_tokens[:valid_image_nums].unsqueeze(-1),
 ]
 )
 return input_ids, image_bounds

 
 @property
 # Copied from transformers.models.clip.processing_clip.CLIPProcessor.model_input_names
 def model_input_names(self):
 tokenizer_input_names = self.tokenizer.model_input_names
 image_processor_input_names = self.image_processor.model_input_names
 return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))


 def pad(self, inputs, max_length=None, padding_value=0, padding_side="left"):
 items = []
 if isinstance(inputs[0], list):
 assert isinstance(inputs[0][0], torch.Tensor)
 for it in inputs:
 for tr in it:
 items.append(tr)
 else:
 assert isinstance(inputs[0], torch.Tensor)
 items = inputs

 batch_size = len(items)
 shape = items[0].shape
 dim = len(shape)
 assert dim <= 2
 if max_length is None:
 max_length = 0
 max_length = max(max_length, max(item.shape[-1] for item in items))
 min_length = min(item.shape[-1] for item in items)
 dtype = items[0].dtype

 if dim == 0:
 return torch.stack([item for item in items], dim=0), [0]
 elif dim == 1:
 if max_length == min_length:
 return torch.stack([item for item in items], dim=0), [0] * batch_size
 tensor = torch.zeros((batch_size, max_length), dtype=dtype) + padding_value
 else:
 tensor = (
 torch.zeros((batch_size, max_length, shape[-1]), dtype=dtype)
 + padding_value
 )

 padding_length = []
 for i, item in enumerate(items):
 if dim == 1:
 if padding_side == "left":
 tensor[i, -len(item) :] = item.clone()
 else:
 tensor[i, : len(item)] = item.clone()
 elif dim == 2:
 if padding_side == "left":
 tensor[i, -len(item) :, :] = item.clone()
 else:
 tensor[i, : len(item), :] = item.clone()
 padding_length.append(tensor.shape[-1] - len(item))

 return tensor, padding_length