159 lines
6.7 KiB
Python
159 lines
6.7 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
from typing import Any, Dict
|
|
|
|
import torch
|
|
from PIL import Image
|
|
from torchvision import transforms
|
|
|
|
from modelscope.preprocessors.image import load_image
|
|
from modelscope.utils.constant import ModeKeys
|
|
from .base import OfaBasePreprocessor
|
|
|
|
|
|
class OfaVisualQuestionAnsweringPreprocessor(OfaBasePreprocessor):
|
|
r"""
|
|
OFA preprocessor for question answer tasks.
|
|
"""
|
|
|
|
def __init__(self,
|
|
cfg,
|
|
model_dir,
|
|
mode=ModeKeys.INFERENCE,
|
|
*args,
|
|
**kwargs):
|
|
"""preprocess the data
|
|
|
|
Args:
|
|
cfg(modelscope.utils.config.ConfigDict) : model config
|
|
model_dir (str): model path,
|
|
mode: preprocessor mode (model mode)
|
|
"""
|
|
super(OfaVisualQuestionAnsweringPreprocessor,
|
|
self).__init__(cfg, model_dir, mode, *args, **kwargs)
|
|
# Initialize transform
|
|
self.patch_resize_transform = transforms.Compose([
|
|
lambda image: image.convert('RGB'),
|
|
transforms.Resize(
|
|
(self.patch_image_size, self.patch_image_size),
|
|
interpolation=transforms.InterpolationMode.BICUBIC),
|
|
transforms.ToTensor(),
|
|
transforms.Normalize(mean=self.mean, std=self.std),
|
|
])
|
|
|
|
def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
|
|
if self.mode == ModeKeys.TRAIN:
|
|
return self._build_train_sample(data)
|
|
else:
|
|
return self._build_infer_sample(data)
|
|
|
|
def _build_train_sample(self, data: Dict[str, Any]) -> Dict[str, Any]:
|
|
r"""
|
|
Building training samples.
|
|
|
|
step 1. Preprocess the data using the logic of `_build_infer_sample`
|
|
and make sure the label data in the result.
|
|
step 2. Preprocessing the label data to generate `target` and `prev_output_token`.
|
|
- add blank in the front out label data and tokenize it as `target` item.
|
|
- if `prompt_type` is `None`, add the bos token as previous output tokens,
|
|
add eos tokens as target items.
|
|
- if `prompt_type` is `src`, concatenate source text input with target item as
|
|
previous output tokens, remove the bos token and add eos token as target items.
|
|
- if `prompt_type` is `prev_output`, just like the `prompt_type` is src, the
|
|
difference is that it will remove the eos token in source text input in this
|
|
setting.
|
|
- padding the source item as final target item.
|
|
step 3. Add constraint mask.
|
|
|
|
Args:
|
|
data (`Dict[str, Any]`): Input data, should contains the key of `image`
|
|
`text` and `label`.
|
|
Return:
|
|
A dict object, contains source text input, patch images, patch masks
|
|
with `Tensor([True])`, decoder prompt, label, target previous output tokens
|
|
and constraint mask.
|
|
"""
|
|
sample = self._build_infer_sample(data)
|
|
tgt_item = self.tokenize_text(
|
|
' {}'.format(sample['label']), add_bos=False, add_eos=False)
|
|
|
|
if self.prompt_type == 'none':
|
|
prev_output_item = torch.cat([self.bos_item, tgt_item])
|
|
target_item = torch.cat([prev_output_item[1:], self.eos_item])
|
|
elif self.prompt_type == 'src':
|
|
prev_output_item = torch.cat([sample['source'], tgt_item])
|
|
target_item = torch.cat([prev_output_item[1:], self.eos_item])
|
|
elif self.prompt_type == 'prev_output':
|
|
prev_output_item = torch.cat([sample['source'][:-1], tgt_item])
|
|
target_item = torch.cat([prev_output_item[1:], self.eos_item])
|
|
else:
|
|
raise NotImplementedError
|
|
target_item[:-len(tgt_item) - 1] = self.tokenizer.pad_token_id
|
|
|
|
sample['prev_output_tokens'] = prev_output_item
|
|
sample['target'] = target_item
|
|
|
|
if self.constraint_trie is not None:
|
|
constraint_mask = torch.zeros(
|
|
(len(target_item), len(self.tgt_dict))).bool()
|
|
start_idx = len(target_item) - len(tgt_item) - 1
|
|
for i in range(
|
|
len(target_item) - len(tgt_item) - 1, len(target_item)):
|
|
constraint_prefix_token = [
|
|
self.tgt_dict.bos()
|
|
] + target_item[start_idx:i].tolist()
|
|
constraint_nodes = self.constraint_trie.get_next_layer(
|
|
constraint_prefix_token)
|
|
constraint_mask[i][constraint_nodes] = True
|
|
sample['constraint_mask'] = constraint_mask
|
|
|
|
return sample
|
|
|
|
def _build_infer_sample(self, data: Dict[str, Any]) -> Dict[str, Any]:
|
|
r"""
|
|
Building inference samples.
|
|
|
|
step 1. Preprocessing image input for model's image input.
|
|
- get pillow image from data.
|
|
- do some transforms to the pillow image, such as resize, normalize etc.
|
|
step 2. Preprocessing the text input for model's text input.
|
|
- add blank in the front of input text.
|
|
- tokenize the result above as source text input.
|
|
step 3. Calculating the decoder prompt.
|
|
- if `prompt_type` is `None`, using bos token.
|
|
- if `prompt_type` is `src`, using source text input
|
|
- if `prompt_type` is `prev_output`, using source text input without eos token.
|
|
step 4. Whether or not to add label data which refer to an answer to the question
|
|
in this task.
|
|
|
|
Args:
|
|
data (`Dict[str, Any]`): Input data, should contains the key of `image`
|
|
`text`.
|
|
Return:
|
|
A dict object, contains source text input, patch images, patch masks
|
|
with `Tensor([True])`, decoder prompt and label.
|
|
"""
|
|
image = self.get_img_pil(data[self.column_map['image']])
|
|
patch_image = self.patch_resize_transform(image)
|
|
text = data[self.column_map['text']]
|
|
text = self.pre_question(text, self.max_src_length)
|
|
text = text + '?' if not text.endswith('?') else text
|
|
inputs = self.tokenize_text(f' {text}')
|
|
if self.prompt_type == 'none':
|
|
decoder_prompt = self.bos_item
|
|
elif self.prompt_type == 'src':
|
|
decoder_prompt = inputs
|
|
elif self.prompt_type == 'prev_output':
|
|
decoder_prompt = inputs[:-1]
|
|
else:
|
|
raise NotImplementedError
|
|
sample = {
|
|
'source': inputs,
|
|
'patch_image': patch_image,
|
|
'patch_mask': torch.tensor([True]),
|
|
'decoder_prompt': decoder_prompt,
|
|
}
|
|
if 'answer' in self.column_map and self.column_map['answer'] in data:
|
|
sample['label'] = data[self.column_map['answer']]
|
|
return sample
|