| import torch |
| from PIL import Image |
|
|
| from transformers import Qwen2VLProcessor, AutoProcessor |
| from transformers.models.qwen2_vl.processing_qwen2_vl import Qwen2VLProcessorKwargs |
|
|
|
|
| class QwenVLProcessor: |
| ROLE = ('user', 'assistant') |
|
|
| def __init__(self, max_length=512, pretrained_model_name_or_path=None): |
| self.processor = AutoProcessor.from_pretrained( |
| pretrained_model_name_or_path) |
| self.max_length = max_length |
|
|
| def __getattr__(self, name: str): |
| try: |
| return super().__getattr__(name) |
| except AttributeError: |
| return getattr(self.processor, name) |
|
|
| def build_prompt(self, query, answer, round=0, system=None): |
| messages = [{"role": self.ROLE[0], "content": query}] |
| if round == 0 and system: |
| messages.insert(0, {"role": "system", "content": system}) |
|
|
| if answer is None: |
| query = self.tokenizer.apply_chat_template( |
| messages, tokenize=False, add_generation_prompt=True) |
| else: |
| messages.append({"role": self.ROLE[1], "content": answer}) |
| prompt = self.tokenizer.apply_chat_template( |
| messages, tokenize=False, add_generation_prompt=False) |
| query, answer = prompt.split("<|im_start|>assistant\n") |
| query += "<|im_start|>assistant\n" |
|
|
| return query, answer |
|
|
| def __call__(self, data_dict, **kwargs): |
| conversations = data_dict["conversations"] |
| images = data_dict.get("images", None) |
| videos = data_dict.get("videos", None) |
|
|
| images = data_dict.get("image", None) |
| if images is not None: |
| images = [Image.open(images).convert('RGB')] |
|
|
| output_kwargs = self._merge_kwargs( |
| Qwen2VLProcessorKwargs, |
| tokenizer_init_kwargs=self.tokenizer.init_kwargs, |
| **kwargs, |
| ) |
| if images is not None: |
| image_inputs = self.image_processor( |
| images=images, videos=None, **output_kwargs["images_kwargs"]) |
| image_grid_thw = image_inputs["image_grid_thw"] |
| else: |
| image_inputs = {} |
| image_grid_thw = None |
|
|
| new_conversation = [] |
| index = 0 |
| for msg in conversations: |
| if msg['from'] == 'human': |
| if image_grid_thw is not None: |
| merge_length = self.image_processor.merge_size**2 |
| text = msg['value'] |
| while "<image>" in text: |
| text = text.replace( |
| "<image>", "<|placeholder|>" * (image_grid_thw[index].prod() // merge_length), 1) |
| index += 1 |
| text = text.replace("<|placeholder|>", "<|image_pad|>") |
| msg['value'] = text |
| new_conversation.append(msg) |
|
|
| input_ids, labels = [], [] |
| for i in range(0, len(new_conversation), 2): |
| query = new_conversation[i]['value'] |
| answer = new_conversation[i+1]['value'] if i + \ |
| 1 < len(new_conversation) else None |
| query, answer = self.build_prompt(query, answer, round=i // 2) |
|
|
| input_ids_ = self.tokenizer( |
| query, add_special_tokens=True, return_attention_mask=False)['input_ids'] |
| labels_ = [-100] * len(input_ids_) |
| if answer is not None: |
| output_ids_ = self.tokenizer(answer, add_special_tokens=True, |
| return_attention_mask=False)['input_ids'] |
| labels_ += output_ids_ |
| input_ids_ += output_ids_ |
| input_ids += input_ids_ |
| labels += labels_ |
|
|
| return { |
| "input_ids": input_ids, |
| "labels": labels, |
| 'pixel_values': image_inputs.get('pixel_values', None), |
| } |
|
|
