"""Datamodule for Llava Pretraining and Finetuning"""
import os
import re
from PIL import Image
import numpy as np
import re
import tempfile
from typing import Dict, List, Union, Tuple
import traceback
import json
import torch
import torch.nn.functional as F
from transformers import DataCollatorForSeq2Seq
from tools.rw_utils import read_jsonlines
from torch.utils.data import Dataset, DataLoader
np_str_obj_array_pattern = re.compile(r"[SaUO]")
default_collate_err_msg_format = (
"default_collate: batch must contain tensors, numpy arrays, numbers, "
"dicts or lists; found {}"
)
from .custom_data_parsers.standard_vision_parser import VisionParser
from .custom_data_parsers.object_tracking_parser import ObjectTrackingParser
from .custom_data_parsers.multi_images_parser import MultiImagesParser
from .custom_data_parsers.video_permutation_parser import VideoPermutationParser
from .custom_data_parsers.utils_visualize import visualize_image_bbox
from .tarsier_processor import TarsierProcessor
from tools.rw_utils import NumpyArrayEncoder
from .utils import DictToObject
import os
HF_TOKEN = os.environ.get('HF_TOKEN', '')
class TarsierDataProcessor:
def __init__(
self,
processor: TarsierProcessor,
n_frames: Union[int, list],
max_n_frames=256,
max_pixels=int(1280 * 720 // 2),
min_pixels=0,
max_seq_len=None,
is_training=True, # 会影响：1. 训练和测试时采帧不同；2. 测试时忽略 response。
print_data_error=True,
do_image_padding=False,
do_image_crop=False,
do_image_resize=True,
video_sampling_strategy={},
prompt='',
train_task='sft',
**kwargs
):
self.kwargs = kwargs
self.processor = processor
self.pad_collator = DataCollatorForSeq2Seq(processor.tokenizer, padding='longest')
self.processor.max_seq_len = self.tokenizer.model_max_length if max_seq_len is None else max_seq_len
self.n_frames = n_frames
self.max_n_frames = max_n_frames
self.max_pixels = max_pixels
self.min_pixels = min_pixels
self.is_training = is_training
self.print_data_error = print_data_error
self.do_image_padding = do_image_padding
self.do_image_crop = do_image_crop
self.do_image_resize = do_image_resize
self.video_sampling_strategy = video_sampling_strategy
self.prompt = prompt
self.train_task = train_task
self.object_tracking_parser = ObjectTrackingParser(
n_frames=self.n_frames,
max_objects=4,
is_training=self.is_training,
)
self.multi_images_parser = MultiImagesParser(
n_frames=self.n_frames,
is_training=self.is_training,
)
self.video_permutation_parser = VideoPermutationParser(
n_frames=self.n_frames,
is_training=self.is_training,
video_sampling_strategy=self.video_sampling_strategy,
)
self.vision_parser = VisionParser(
n_frames=self.n_frames,
max_n_frames=self.max_n_frames,
is_training=self.is_training,
video_sampling_strategy=self.video_sampling_strategy
)
def select_parser(self, data_dict):
if data_dict.get('task', None) == 'video/object_tracking':
return self.object_tracking_parser
elif data_dict.get('task', None) == 'multi_images':
return self.multi_images_parser
elif data_dict.get('dataset', None) == 'video_permutation':
return self.video_permutation_parser
else:
return self.vision_parser
def parse_image_processing_config(self, data_dict):
image_processing_config=data_dict.get('image_processing_config', {})
do_padding = image_processing_config.get('do_padding', self.do_image_padding)
do_crop = image_processing_config.get('do_crop', self.do_image_crop)
do_resize = image_processing_config.get('do_resize', self.do_image_resize)
max_pixels = image_processing_config.get('max_pixels', self.max_pixels)
min_pixels = image_processing_config.get('min_pixels', self.min_pixels)
assert min_pixels <= max_pixels
image_processing_config['do_padding'] = do_padding
image_processing_config['do_crop'] = do_crop
image_processing_config['do_resize'] = do_resize
image_processing_config['max_pixels'] = max_pixels
image_processing_config['min_pixels'] = min_pixels
return image_processing_config
def _transform(self, raw_data_dict: Dict) -> Dict:
data_dict = json.loads(json.dumps(raw_data_dict, cls=NumpyArrayEncoder))
del raw_data_dict
if self.prompt:
for msg in data_dict['messages']:
if msg['role'] == 'user':
for content in msg['content']:
if content['type'] == 'text':
content['text'] = self.prompt
data_dict_copy = json.loads(json.dumps(data_dict, cls=NumpyArrayEncoder))
image_processing_config = self.parse_image_processing_config(data_dict)
parser = self.select_parser(data_dict)
messages = parser.transform(data_dict, image_processing_config)
data_dict_copy['extra_info'] = data_dict.pop('extra_info', {})
# visualize_image_bbox(data_dict, image_processing_config, self.processor)
outputs = self.processor(messages, image_processing_config, is_training=self.is_training)
# if not self.is_training:
outputs['raw_data_dict'] = data_dict_copy
return [outputs]
def _split_chosen_rejected(self, data_dict: Dict):
chosen_data_dict = data_dict
rejected_data_dict = json.loads(json.dumps(data_dict, cls=NumpyArrayEncoder))
for msg in chosen_data_dict['messages']:
if msg['role'] == 'assistant':
for content in msg['content']:
if content['type'] == 'text':
content['text'] = content['chosen']
for msg in rejected_data_dict['messages']:
if msg['role'] == 'assistant':
for content in msg['content']:
if content['type'] == 'text':
content['text'] = content['rejected']
return chosen_data_dict, rejected_data_dict
def transform(self, data_dict: Dict) -> Dict:
try:
if self.train_task == 'dpo':
chosen_data_dict, rejected_data_dict = self._split_chosen_rejected(data_dict)
return self._transform(chosen_data_dict) + self._transform(rejected_data_dict)
return self._transform(data_dict)
except Exception as e:
if self.print_data_error:
print(traceback.format_exc())
print(f'Error occurs when processing: \n{data_dict}')
return []
def batch_transform(self, batch_data: List[Dict]) -> Dict:
model_inputs = {}
# if not self.is_training:
raw_data_dict = [d.pop('raw_data_dict') for d in batch_data]
model_inputs['raw_data_dict'] = raw_data_dict
batch_pixel_values = [d.pop('pixel_values') for d in batch_data if 'pixel_values' in d]
batch_image_grid_thw = [d.pop('image_grid_thw') for d in batch_data if 'image_grid_thw' in d]
if len(batch_pixel_values) == 0:
vision_placeholder = self.get_vision_placeholder()
batch_pixel_values = [vision_placeholder.get('pixel_values')]
batch_image_grid_thw = [vision_placeholder.get('image_grid_thw')] if 'image_grid_thw' in vision_placeholder else []
model_inputs['pixel_values'] = torch.cat(batch_pixel_values, dim=0)
if len(batch_image_grid_thw) > 0:
model_inputs['image_grid_thw'] = torch.cat(batch_image_grid_thw, dim=0)
batch_num_images = [d.pop('num_images') for d in batch_data]
model_inputs['num_images'] = torch.tensor(batch_num_images)
model_inputs.update(self.pad_collator(batch_data))
return model_inputs
def __call__(self, batch_data: Union[Dict, List[Dict]]) -> Dict:
if isinstance(batch_data, dict):
batch_data = [batch_data]
batch = [self.transform(d)[0] for d in batch_data]
return self.batch_transform(batch)
def get_vision_placeholder(self):
messages = [{"role": "user", "content": [{"type": "image", "image": Image.new(mode='RGB', size=(336, 336))}]}]
image_processing_config = self.parse_image_processing_config({})
return self.processor(messages, image_processing_config)
def get_text_placeholder(self):
messages = [
{"role": "user", "content": [{"type": "text", "text": "Hello!"}]},
{"role": "assistant", "content": [{"type": "text", "text": "Thank you very much"}]},
]
image_processing_config = self.parse_image_processing_config({})
return self.processor(messages, image_processing_config)
def init_processor(processor: Union[TarsierProcessor, str]=None, config: Dict=None):
config = DictToObject(config) if isinstance(config, dict) else config
if isinstance(processor, str):
sub_processor = TarsierProcessor.from_pretrained(
processor,
padding_side='left',
trust_remote_code=True,
token=HF_TOKEN,
)
else:
sub_processor = processor
processor = TarsierDataProcessor(
processor=sub_processor,
n_frames=config.n_frames,
max_n_frames=config.max_n_frames,
max_pixels=config.max_pixels,
min_pixels=config.min_pixels,
max_seq_len=config.max_seq_len,
is_training=config.is_training,
print_data_error=config.print_data_error,
do_image_padding=config.do_image_padding,
do_image_crop=config.do_image_crop,
do_image_resize=config.do_image_resize,
video_sampling_strategy=config.video_sampling_strategy,
prompt=config.prompt,
train_task=config.train_task
)
return processor
class TarsierDataset(Dataset):
def __init__(self, ann_path="", anns=None, config: Dict=None, processor: Union[TarsierDataProcessor, TarsierProcessor, str]=None):
self.config = DictToObject(config) if isinstance(config, dict) else config
if not isinstance(processor, TarsierDataProcessor):
self.processor = init_processor(processor, config)
else:
self.processor = processor
if anns is None:
self.anns = []
if isinstance(ann_path, str):
ann_path = [ann_path]
for path in ann_path:
self.anns.extend(read_jsonlines(path))
else:
self.anns = anns
def __len__(self):
return len(self.anns)
def __getitem__(self, index):
if index < 0 or index >= len(self.anns):
raise IndexError("Index out of range")
try:
ann = self.anns[index]
model_inputs = self.processor(ann)
except Exception as e:
print(f"Load data error: {e}")
return ann, None
return ann, model_inputs