training source (data.py, train.py, ResumeDatasetCallback)
1c32e03 verified 4 months ago
35.7 kB
import copy
import os
from dataclasses import dataclass, field
from typing import Dict
import torch
import transformers
import ujson as json
from torch.utils.data import Dataset
from qwen_vl_utils import process_vision_info
from PIL import Image
from transformers import AutoImageProcessor
import re
import numpy as np
import cv2
from torchvision import transforms
import random

from segment_anything import build_sam_vit_h, sam_model_registry, SamPredictor
from src.anchors.DepthAnything.depth_anything_v2.dpt import DepthAnythingV2
from diffusers import AutoencoderKL
from transformers import AutoModel, CLIPImageProcessor

from .params import DataArguments
from .constants import *


def truncate_sequence(input_ids, labels, max_length, eos_token_id):
    if input_ids.size(0) > max_length:
        input_ids = input_ids[:max_length-1]
        labels = labels[:max_length-1]

    if eos_token_id is not None:
        input_ids = torch.cat([input_ids, torch.tensor([eos_token_id])])
        labels = torch.cat([labels, torch.tensor([eos_token_id])])

    return input_ids, labels

def pad_sequence(sequences, padding_side='right', padding_value=0):
    """
    Pad a list of sequences to the same length.
    sequences: list of tensors in [seq_len, *] shape
    """
    assert padding_side in ['right', 'left']
    max_size = sequences[0].size()
    trailing_dims = max_size[1:]
    max_len = max(len(seq) for seq in sequences)
    batch_size = len(sequences)
    output = sequences[0].new_full((batch_size, max_len) + trailing_dims, padding_value)
    for i, seq in enumerate(sequences):
        length = seq.size(0)
        if padding_side == 'right':
            output.data[i, :length] = seq
        else:
            output.data[i, -length:] = seq
    return output

def get_image_info(image_path, min_pixel, max_pixel, width, height):
    # Using this because of process_vision_info function
    # Need to fix this in the future


    content = {
        "type": "image", 
        "image": image_path,
        "min_pixel": min_pixel,
        "max_pixel": max_pixel
    }

    if width is not None and height is not None:
        content["resized_width"] = width
        content["resized_height"] = height

    messages = [
        {"role": "user", 
         "content": [content]
        }
    ]

    image_input, _ = process_vision_info(messages)

    return image_input[0]

def get_video_info(video_path, min_pixels, max_pixels, fps):
    # Using this because of process_vision_info function
    # Need to fix this in the future

    messages = [
        {"role": "user", 
         "content": [
             {
                "type": "video", 
                "video": video_path,
                "min_pixels": min_pixels,
                "max_pixels": max_pixels,
                "fps": fps
            }
            ]
        }
    ]

    _, video_input, video_kwargs = process_vision_info(messages, return_video_kwargs=True)

    return video_input[0], video_kwargs

def add_anchor_pad(user_input, anchor_nums, anchor_tokens):
    # add anchor pad after VISION_END_TOKEN or ANCHOR_END_TOKEN
    anchor_pads = []
    for anchor_num, anchor_token in zip(anchor_nums, anchor_tokens):
        anchor_pad = ANCHOR_START_TOKEN + anchor_token * anchor_num + ANCHOR_END_TOKEN
        anchor_pads.append(anchor_pad)
    anchor_pads = "".join(anchor_pads)
    if VISION_END_TOKEN in user_input:
        user_input = user_input.replace(VISION_END_TOKEN, VISION_END_TOKEN + anchor_pads)
    return user_input

def add_cot_anchor_pad_in_user_input(user_input, anchor_nums, anchor_tokens):
    if len(anchor_nums) == 0:
        return user_input
    
    anchor_pads = []
    for anchor_num, anchor_token in zip(anchor_nums, anchor_tokens, anchor_names):
        anchor_pad = ANCHOR_START_TOKEN + anchor_token * anchor_num + ANCHOR_END_TOKEN
        anchor_pads.append(anchor_pad)
    CoT_pad = ""
    if len(anchor_pads) == 1:
        CoT_pad = f"The {anchor_names[0]} of the image is {anchor_pads[0]}. "
    else:
        for i, (anchor_name, anchor_pad) in enumerate(zip(anchor_names, anchor_pads)):
            if i == 0:
                CoT_pad += f"The {anchor_name} of the image is {anchor_pad}, "
            elif i == len(anchor_names) - 1:
                CoT_pad += f"and the {anchor_name} of the image is {anchor_pad}. "
            else:
                CoT_pad += f"the {anchor_name} of the image is {anchor_pad}, "
    user_input = CoT_pad + user_input
    return user_input

def get_cot_data_in_response(response, anchor_nums, anchor_tokens, anchor_names):
    if len(anchor_nums) == 0:
        return response
    
    anchor_pads = []
    for anchor_num, anchor_token in zip(anchor_nums, anchor_tokens):
        anchor_pad = ANCHOR_START_TOKEN + anchor_token * anchor_num + ANCHOR_END_TOKEN
        anchor_pads.append(anchor_pad)
    CoT_start = "Because "
    if len(anchor_names) == 1:
        CoT_start += f"the {anchor_names[0]} of the image is {anchor_pads[0]}. "
    else:
        for anchor_name, anchor_pad in zip(anchor_names, anchor_pads):
            CoT_start += f"the {anchor_name} of the image is {anchor_pad}"
            if anchor_name == anchor_names[-2]:
                CoT_start += ", and "
            elif anchor_name == anchor_names[-1]:
                CoT_start += ". "
            else:
                CoT_start += ", "
    response = CoT_start + response
    return response


COT_TEMPLATES = [
    {
        "name": "basic_causal",
        "single": "Because the {anchor_name} of the image is {anchor_pad}. ",
        "multiple": "Because the {anchor_name} of the image is {anchor_pad}{connector}",
        "connectors": {
            "middle": ", ",
            "second_last": ", and ",
            "last": ". "
        }
    },
    
    {
        "name": "observational",
        "single": "I can observe that the {anchor_name} of the image is {anchor_pad}. ",
        "multiple": "I can observe that the {anchor_name} of the image is {anchor_pad}{connector}",
        "connectors": {
            "middle": ", ",
            "second_last": ", and ",
            "last": ". "
        }
    },
    
    {
        "name": "analytical",
        "single": "After analyzing the image, the {anchor_name} is {anchor_pad}. ",
        "multiple": "After analyzing the image, the {anchor_name} is {anchor_pad}{connector}",
        "connectors": {
            "middle": ", ",
            "second_last": ", and ",
            "last": ". "
        }
    },
    
    {
        "name": "descriptive",
        "single": "The image shows that the {anchor_name} is {anchor_pad}. ",
        "multiple": "The image shows that the {anchor_name} is {anchor_pad}{connector}",
        "connectors": {
            "middle": ", ",
            "second_last": ", and ",
            "last": ". "
        }
    },
    
    {
        "name": "conditional",
        "single": "Given that the {anchor_name} of the image is {anchor_pad}. ",
        "multiple": "Given that the {anchor_name} of the image is {anchor_pad}{connector}",
        "connectors": {
            "middle": ", ",
            "second_last": ", and ",
            "last": ". "
        }
    },
    
    {
        "name": "evidence_based",
        "single": "Based on the visual evidence, the {anchor_name} of the image is {anchor_pad}. ",
        "multiple": "Based on the visual evidence, the {anchor_name} of the image is {anchor_pad}{connector}",
        "connectors": {
            "middle": ", ",
            "second_last": ", and ",
            "last": ". "
        }
    }
]

def get_random_cot_template():
    return random.choice(COT_TEMPLATES)

def apply_cot_template(template, anchor_names, anchor_pads):
    if len(anchor_names) == 1:
        return template["single"].format(
            anchor_name=anchor_names[0], 
            anchor_pad=anchor_pads[0]
        )
    else:
        result = ""
        for i, (anchor_name, anchor_pad) in enumerate(zip(anchor_names, anchor_pads)):
            if i == len(anchor_names) - 1: 
                connector = template["connectors"]["last"]
            elif i == len(anchor_names) - 2: 
                connector = template["connectors"]["second_last"]
            else: 
                connector = template["connectors"]["middle"]
            
            result += template["multiple"].format(
                anchor_name=anchor_name,
                anchor_pad=anchor_pad,
                connector=connector
            )
        return result

def get_templates_comt_data_in_response(response, anchor_nums, anchor_tokens, anchor_names):
    if len(anchor_nums) == 0:
        return response
    
    anchor_pads = []
    for anchor_num, anchor_token in zip(anchor_nums, anchor_tokens):
        anchor_pad = ANCHOR_START_TOKEN + anchor_token * anchor_num + ANCHOR_END_TOKEN
        anchor_pads.append(anchor_pad)
    
    template = get_random_cot_template()
    
    cot_text = apply_cot_template(template, anchor_names, anchor_pads)
    
    response = "<think>" + cot_text + "</think>" + "<answer>" + response + "</answer>"
    return response
        
def get_comt_data_in_response(response, anchor_nums, anchor_tokens, anchor_names):
    if len(anchor_nums) == 0:
        return response
    
    anchor_pads = []
    for anchor_num, anchor_token in zip(anchor_nums, anchor_tokens):
        anchor_pad = ANCHOR_START_TOKEN + anchor_token * anchor_num + ANCHOR_END_TOKEN
        anchor_pads.append(anchor_pad)
    CoT_start = "<think> Because "
    if len(anchor_names) == 1:
        CoT_start += f"the {anchor_names[0]} of the image is {anchor_pads[0]}. "
    else:
        for anchor_name, anchor_pad in zip(anchor_names, anchor_pads):
            CoT_start += f"the {anchor_name} of the image is {anchor_pad}"
            if anchor_name == anchor_names[-2]:
                CoT_start += ", and "
            elif anchor_name == anchor_names[-1]:
                CoT_start += ". "
            else:
                CoT_start += ", "
    response = CoT_start + " </think>\n" + "<answer> " + response + " </answer>"
    return response
    
def get_feature_data(user_input, gpt_response, anchor_nums, anchor_tokens, anchor_names):
    anchor_pads = []
    for anchor_num, anchor_token, anchor_name in zip(anchor_nums, anchor_tokens, anchor_names):
        anchor_pad = ANCHOR_START_TOKEN + anchor_token * anchor_num + ANCHOR_END_TOKEN
        anchor_pads.append(anchor_pad)
    anchor_name = ", ".join(anchor_names)
    anchor_pads = "".join(anchor_pads)
    user_input = f"{DEFAULT_IM_START_TOKEN}{user_input['role']}\n{VISION_START_TOKEN + DEFAULT_IMAGE_TOKEN + VISION_END_TOKEN}What is the {anchor_name} of the image?\n{DEFAULT_IM_END_TOKEN}\n{DEFAULT_IM_START_TOKEN}{gpt_response['role']}\n"
    gpt_response = f"{anchor_pads}\n{DEFAULT_IM_END_TOKEN}\n"
    return user_input, gpt_response

def replace_pad_with_anchor_tokens(gpt_response):
    token_dict = {
        "<segmentation>": SAM_PAD_TOKEN * 8,
        "<depth>": DEPTH_PAD_TOKEN * 4,
        "<dino>": DINO_PAD_TOKEN * 4,
        "<pidinet>": PIDINET_PAD_TOKEN * 4,
        "<siglip>": SIGLIP_PAD_TOKEN * 4,
        "<metaclip>": METACLIP_PAD_TOKEN * 4,
    }
    for token, anchor_token in token_dict.items():
        gpt_response = gpt_response.replace(token, anchor_token)
    return gpt_response
    
def get_token_num(anchor_model_id):
    token_nums = []
    for anchor_model in anchor_model_id:
        if anchor_model == "sam":
            token_nums.append(8)
        elif anchor_model == "dino":
            token_nums.append(4)
        elif anchor_model == "depth":
            token_nums.append(4)
        elif anchor_model == "SD":
            token_nums.append(4)
        elif anchor_model == "InternViT":
            token_nums.append(4)
        elif anchor_model == "pidinet":
            token_nums.append(4)
        elif anchor_model == "siglip":
            token_nums.append(4)
        elif anchor_model == "metaclip":
            token_nums.append(4)
    return token_nums

def get_anchor_token(anchor_model_id):
    anchor_tokens = []
    for anchor_model in anchor_model_id:
        if anchor_model == "sam":
            anchor_tokens.append(SAM_PAD_TOKEN)
        elif anchor_model == "dino":
            anchor_tokens.append(DINO_PAD_TOKEN)
        elif anchor_model == "depth":
            anchor_tokens.append(DEPTH_PAD_TOKEN)
        elif anchor_model == "SD":
            anchor_tokens.append(SD_PAD_TOKEN)
        elif anchor_model == "InternViT":
            anchor_tokens.append(INTERN_PAD_TOKEN)
        elif anchor_model == "pidinet":
            anchor_tokens.append(PIDINET_PAD_TOKEN)
        elif anchor_model == "siglip":
            anchor_tokens.append(SIGLIP_PAD_TOKEN)
        elif anchor_model == "metaclip":
            anchor_tokens.append(METACLIP_PAD_TOKEN)
    return anchor_tokens

def get_anchor_task_name(anchor_model_id):
    anchor_task_names = []
    for anchor_model in anchor_model_id:
        if anchor_model == "sam":
            anchor_task_names.append("segmentation")
        elif anchor_model == "dino":
            anchor_task_names.append("perception feature")
        elif anchor_model == "depth":
            anchor_task_names.append("depth map")
        elif anchor_model == "SD":
            anchor_task_names.append("style")
        elif anchor_model == "InternViT":
            anchor_task_names.append("caption")
        elif anchor_model == "pidinet":
            anchor_task_names.append("edge map")
        elif anchor_model == "siglip":
            anchor_task_names.append("clip feature")
        elif anchor_model == "metaclip":
            anchor_task_names.append("metaclip feature")
    return anchor_task_names
    

class SupervisedDataset(Dataset):
    """Dataset for supervised fine-tuning."""

    def __init__(
        self,
        data_path: str | list,
        processor: transformers.ProcessorMixin,
        data_args: DataArguments,
        model_id,
        padding=True,
        shuffle=True,
        random_seed=42,
        anchor_model_id=None,
    ):
        super(SupervisedDataset, self).__init__()
        if isinstance(data_path, str):
            import os as _os, glob as _glob
            if _os.path.isdir(data_path):
                # Parquet streaming: no HF cache, reads rows on-demand via pyarrow
                import pyarrow.parquet as _pq
                import bisect as _bisect
                parquet_files = sorted(_glob.glob(_os.path.join(data_path, "**", "*.parquet"), recursive=True))
                if not parquet_files:
                    raise FileNotFoundError(f"No parquet files found under {data_path}")
                print(f"[DataLoader] Found {len(parquet_files)} parquet files (streaming mode)")
                cumulative_rows = []
                cumsum = 0
                for _f in parquet_files:
                    _pf = _pq.ParquetFile(_f)
                    cumsum += _pf.metadata.num_rows
                    cumulative_rows.append(cumsum)
                self._is_parquet = True
                self._parquet_files = parquet_files
                self._cumulative_rows = cumulative_rows
                self._total_rows = cumsum
                print(f"[DataLoader] Total parquet rows: {cumsum}")
                self._pf_handles = {}
                self._row_group_index = {}
                list_data_dict = None
            else:
                self._is_parquet = False
                list_data_dict = json.load(open(data_path, "r"))
        else:
            self._is_parquet = False
            list_data_dict = data_path

        self.model_id = model_id
        self.processor = processor
        self.list_data_dict = list_data_dict
        self.data_args = data_args
        self.padding = padding
        self.image_min_pixel = data_args.image_min_pixels
        self.image_max_pixel = data_args.image_max_pixels
        self.image_resized_w = data_args.image_resized_width
        self.image_resized_h = data_args.image_resized_height
        self.video_min_pixel = data_args.video_min_pixels
        self.video_max_pixel = data_args.video_max_pixels
        self.fps = data_args.fps
        self.anchor_model_id = anchor_model_id
        self.anchor_token_nums = get_token_num(anchor_model_id)
        self.anchor_tokens = get_anchor_token(anchor_model_id)
        self.anchor_task_names = get_anchor_task_name(anchor_model_id)
        
        self.cur_step = 0
        self.stage_0_step = data_args.stage_0_step
        self.stage_1_step = data_args.stage_1_step
        self.stage_2_step = data_args.stage_2_step
        
        # for shuffle
        self.rng = np.random.default_rng(seed=random_seed)
        
        if shuffle:
            if self._is_parquet:
                import numpy as _np
                self._shuffle_perm = _np.random.RandomState(random_seed).permutation(self._total_rows)
            else:
                from datasets import Dataset as _HFDataset
                if isinstance(self.list_data_dict, _HFDataset):
                    self.list_data_dict = self.list_data_dict.shuffle(seed=random_seed)
                else:
                    self.rng.shuffle(self.list_data_dict)
        else:
            if self._is_parquet:
                self._shuffle_perm = None
    
    def set_cur_step(self, step: int):
        self.cur_step = step
        print(f"[Dataset] cur_step has been set to {step}")
        
    def __len__(self):
        if self._is_parquet:
            return self._total_rows
        return len(self.list_data_dict)

    def _get_parquet_row(self, idx):
        import pyarrow.parquet as pq, bisect, io
        real_idx = self._shuffle_perm[idx] if self._shuffle_perm is not None else idx
        # bisect to find file
        fi = bisect.bisect_right(self._cumulative_rows, real_idx)
        local_idx = real_idx - self._cumulative_rows[fi - 1] if fi > 0 else real_idx
        fpath = self._parquet_files[fi]

        # Open ParquetFile lazily and cache handle
        if fpath not in self._pf_handles:
            self._pf_handles[fpath] = pq.ParquetFile(fpath)
            pf = self._pf_handles[fpath]
            # Build row-group cumulative table
            cum = []
            s = 0
            for rg_idx in range(pf.num_row_groups):
                s += pf.metadata.row_group(rg_idx).num_rows
                cum.append(s)
            self._row_group_index[fpath] = cum

        pf = self._pf_handles[fpath]
        cum = self._row_group_index[fpath]

        # Find which row group contains local_idx
        rg_idx = bisect.bisect_right(cum, local_idx)
        rg_start = 0 if rg_idx == 0 else cum[rg_idx - 1]
        in_rg_idx = local_idx - rg_start

        # Read just this row group (typically ~1k-10k rows, much smaller than file)
        table = pf.read_row_group(rg_idx)
        row = table.slice(in_rg_idx, 1).to_pylist()[0]

        # Image decode
        image_data = row.get('image')
        if isinstance(image_data, dict) and image_data.get('bytes'):
            row['image'] = Image.open(io.BytesIO(image_data['bytes'])).convert('RGB')
        elif isinstance(image_data, dict) and image_data.get("path"):
            row["image"] = Image.open(image_data["path"]).convert("RGB")

        return row
    def __getitem__(self, i) -> Dict[str, torch.Tensor]:
        
        # import ipdb; ipdb.set_trace()
        
        if self._is_parquet:
            sources = self._get_parquet_row(i)
        else:
            sources = self.list_data_dict[i]
        
        is_video = False

        processor = self.processor
        if "image" in sources:
            videos = None
            grid_key = "image_grid_thw"
            pixel_key = "pixel_values"
            
            image_files = sources["image"]
            image_folder = self.data_args.image_folder

            if isinstance(image_files, Image.Image):
                # Already a PIL Image (e.g. from parquet HF dataset)
                image_files = [image_files.convert("RGB")]
            elif isinstance(image_files, str):
                image_files = Image.open(image_files).convert("RGB")
                image_files = [image_files]
            elif isinstance(image_files, (bytes, bytearray)):
                import io as _io
                image_files = [Image.open(_io.BytesIO(image_files)).convert("RGB")]
            else:
                image_files = [img_f.convert("RGB") if isinstance(img_f, Image.Image) else Image.open(img_f).convert("RGB") for img_f in image_files]

            images = []
            
            for image_file in image_files:
                
                # if not os.path.exists(image_file):
                #     if not image_file.startswith("http"):
                #         image_file = os.path.join(image_folder, image_file)
                #         images.append(get_image_info(image_file, self.image_min_pixel, self.image_max_pixel, self.image_resized_w, self.image_resized_h))

                # else:
                images.append(get_image_info(image_file, self.image_min_pixel, self.image_max_pixel, self.image_resized_w, self.image_resized_h))

        elif "video" in sources:
            is_video = True
            images=None
            grid_key = "video_grid_thw"
            pixel_key = "pixel_values_videos"

            video_files = sources["video"]
            video_folder = self.data_args.image_folder

            if isinstance(video_files, str):
                video_files = [video_files]

            videos = []
            for video_file in video_files:
                if not os.path.exists(video_file):
                    if not video_file.startswith("http"):
                        video_file = os.path.join(video_folder, video_file)
                video_input, video_kwargs = get_video_info(video_file, self.video_min_pixel, self.video_max_pixel, self.data_args.fps)
                videos.append(video_input)
        else:
            grid_key = None
            pixel_key = None
            images=None
            videos=None
           
        if images is None:
            
            print("No image or video found in the data.")
            images = []
            # Create a black image as a placeholder
            black_image = Image.new("RGB", (self.image_resized_w, self.image_resized_h), (0, 0, 0))
            images.append(get_image_info(black_image, self.image_min_pixel, self.image_max_pixel, self.image_resized_w, self.image_resized_h))

        elif len(images) == 0:
            print("No image or video found in the data.")
            # Create a black image as a placeholder
            black_image = Image.new("RGB", (self.image_resized_w, self.image_resized_h), (0, 0, 0))
            images.append(get_image_info(black_image, self.image_min_pixel, self.image_max_pixel, self.image_resized_w, self.image_resized_h))
        
        if videos is not None:
            
            # import ipdb; ipdb.set_trace()
            pass
            
        sources = copy.deepcopy(llava_to_openai(sources['conversations'], is_video=is_video))

        all_input_ids = [] 
        all_labels = []
        all_pixel_values = []
        all_image_grid_thw = []
        all_second_gird = []
        
        # all_dino_encoded_values = []

        # Qwen2-VL uses a default system message so I've added this.
        if len(SYSTEM_MESSAGE) > 0:
            system_message = f"{DEFAULT_IM_START_TOKEN}system\n{SYSTEM_MESSAGE}\n{DEFAULT_IM_END_TOKEN}\n"
            system_message_input_ids = processor.tokenizer(system_message, add_special_tokens=False, return_tensors='pt')['input_ids']
            system_labels = torch.full_like(system_message_input_ids, IGNORE_INDEX) 
            
            all_input_ids.append(system_message_input_ids.squeeze(0))
            all_labels.append(system_labels.squeeze(0))
            
        for _, j in enumerate(range(0, len(sources), 2)):
            
            if j >= 2:
                break
            
            user_input = sources[j]
            gpt_response = sources[j + 1]
                        
            if (DEFAULT_IMAGE_TOKEN not in user_input['content']) and (DEFAULT_VIDEO_TOKEN not in user_input['content']) and (LLAVA_IMAGE_TOKEN in user_input['content']):
                user_input = f"{DEFAULT_IM_START_TOKEN}{VISION_START_TOKEN + DEFAULT_IMAGE_TOKEN + VISION_END_TOKEN}{user_input['role']}\n{user_input['content']}\n{DEFAULT_IM_END_TOKEN}\n{DEFAULT_IM_START_TOKEN}{gpt_response['role']}\n"
                user_input = add_anchor_pad(user_input, self.anchor_token_nums, self.anchor_tokens)
                gpt_response = f"{gpt_response['content']}\n{DEFAULT_IM_END_TOKEN}\n"
                raise ValueError('Every man is a poet when he is in love')
            else:
                if self.cur_step < self.stage_0_step:
                    user_input = f"{DEFAULT_IM_START_TOKEN}{user_input['role']}\n{user_input['content']}\n{DEFAULT_IM_END_TOKEN}\n{DEFAULT_IM_START_TOKEN}{gpt_response['role']}\n"
                    user_input = add_anchor_pad(user_input, self.anchor_token_nums, self.anchor_tokens)
                    gpt_response = f"{gpt_response['content']}\n{DEFAULT_IM_END_TOKEN}\n"
                elif self.cur_step < self.stage_1_step:
                    user_input, gpt_response = get_feature_data(user_input, gpt_response, self.anchor_token_nums, self.anchor_tokens, self.anchor_task_names)
                elif self.cur_step < self.stage_2_step:
                    user_input = f"{DEFAULT_IM_START_TOKEN}{user_input['role']}\n{user_input['content']}\n{DEFAULT_IM_END_TOKEN}\n{DEFAULT_IM_START_TOKEN}{gpt_response['role']}\n"
                    gpt_response = f"{gpt_response['content']}"
                    if DEFAULT_IMAGE_TOKEN in user_input:
                        gpt_response = get_comt_data_in_response(gpt_response, self.anchor_token_nums, self.anchor_tokens, self.anchor_task_names)
                    gpt_response = f"{gpt_response}\n{DEFAULT_IM_END_TOKEN}\n"
                    # print(f"\033[92m gpt_response: {gpt_response}\033[0m")
                else:
                    # user_input = f"{DEFAULT_IM_START_TOKEN}{user_input['role']}\n{user_input['content']}\n{DEFAULT_IM_END_TOKEN}\n{DEFAULT_IM_START_TOKEN}{gpt_response['role']}\n"
                    # gpt_response = f"{gpt_response['content']}\n{DEFAULT_IM_END_TOKEN}\n"
                    # gpt_response = replace_pad_with_anchor_tokens(gpt_response)
                    import random
                    xxx = random.randint(0, 5)
                    if xxx == 0:
                        user_input = f"{DEFAULT_IM_START_TOKEN}{user_input['role']}\n{user_input['content']}\n{DEFAULT_IM_END_TOKEN}\n{DEFAULT_IM_START_TOKEN}{gpt_response['role']}\n"
                        gpt_response = f"{gpt_response['content']}\n{DEFAULT_IM_END_TOKEN}\n"
                    else:
                        user_input = f"{DEFAULT_IM_START_TOKEN}{user_input['role']}\n{user_input['content']}\n{DEFAULT_IM_END_TOKEN}\n{DEFAULT_IM_START_TOKEN}{gpt_response['role']}\n"
                        gpt_response = f"{gpt_response['content']}"
                        if DEFAULT_IMAGE_TOKEN in user_input:
                            # INSERT_YOUR_CODE
                            total = len(self.anchor_tokens)
                            if total == 0:
                                selected_anchor_token_nums = []
                                selected_anchor_tokens = []
                                selected_anchor_task_names = []
                            else:
                                x = random.randint(1, total)
                                idxs = sorted(random.sample(range(total), x)) if x > 0 else []
                                selected_anchor_token_nums = [self.anchor_token_nums[i] for i in idxs]
                                selected_anchor_tokens = [self.anchor_tokens[i] for i in idxs]
                                selected_anchor_task_names = [self.anchor_task_names[i] for i in idxs]
                            gpt_response = get_comt_data_in_response(gpt_response, selected_anchor_token_nums, selected_anchor_tokens, selected_anchor_task_names)
                        gpt_response = f"{gpt_response}\n{DEFAULT_IM_END_TOKEN}\n"
            
            # print(f'the user_input is {user_input}')
            # print(f'the gpt_response is {gpt_response}')
                    

            # print("-----------------")
            # print(user_input, gpt_response)
            # print("-----------------")
            
            # import ipdb; ipdb.set_trace()
            
            if DEFAULT_IMAGE_TOKEN in user_input:
                inputs = processor(text=[user_input], images=images, videos=videos, padding=False, return_tensors='pt')
                prompt_input_ids = inputs['input_ids']
                # raise ValueError('Every man is a poet when he is in love')
                all_pixel_values.append(inputs[pixel_key])
                all_image_grid_thw.append(inputs[grid_key])
                
                # del dino_val
                torch.cuda.empty_cache()
                
            
            elif DEFAULT_VIDEO_TOKEN in user_input:
                if "Qwen2.5" in self.model_id:
                    inputs = processor(text=[user_input], images=images, videos=videos, padding=False, return_tensors='pt', **video_kwargs)
                    all_second_gird.extend(inputs["second_per_grid_ts"])
                else:
                    inputs = processor(text=[user_input], images=images, videos=videos, padding=False, return_tensors='pt')
                prompt_input_ids = inputs['input_ids']
                all_pixel_values.append(inputs[pixel_key])
                all_image_grid_thw.append(inputs[grid_key])

            else:
                prompt_input_ids = processor.tokenizer(user_input, add_special_tokens=False, padding=False, return_tensors='pt')['input_ids']

            response_input_ids = processor.tokenizer(gpt_response, add_special_tokens=False, padding=False, return_tensors='pt')['input_ids']

            input_ids = torch.cat([prompt_input_ids, response_input_ids], dim=1).squeeze(0)
            labels = torch.cat(
                [
                    torch.tensor([IGNORE_INDEX] * len(prompt_input_ids[0])),  
                    response_input_ids.squeeze(0),
                ],
                dim=0,
            )

            all_input_ids.append(input_ids)
            all_labels.append(labels)
        
        # There is no need for eos or bos tokens in the input_ids
        # Qwen2-VL does not use them
        input_ids = torch.cat(all_input_ids, dim=0).to(torch.long)
        labels = torch.cat(all_labels, dim=0).to(torch.long)

        # eos_token_id = processor.tokenizer.convert_tokens_to_ids(DEFAULT_IM_END_TOKEN)
        # input_ids, labels = truncate_sequence(input_ids, labels, self.max_length, eos_token_id)

        attention_mask = (input_ids > -1000000).to(torch.long)
        
        
        
        data_dict = dict(
            input_ids=input_ids,
            attention_mask=attention_mask,
            labels=labels,
        )

        if pixel_key and grid_key:
            pixel_values = torch.cat(all_pixel_values, dim=0)
            image_thw = torch.cat(all_image_grid_thw, dim=0)
            
            data_dict[pixel_key] = pixel_values
            data_dict[grid_key] = image_thw
            data_dict["image_files"] = image_files

        if len(all_second_gird) > 0:
            second_gird = all_second_gird
            data_dict["second_per_grid_ts"] = second_gird
        
        self.cur_step += 1
        
        return data_dict


class DataCollatorForSupervisedDataset(object):
    """Collate examples for supervised fine-tuning."""

    def __init__(self, pad_token_id: int):
        self.pad_token_id = pad_token_id

    def __call__(self, examples):
        batch_input_ids = []
        batch_label_ids = []
        batch_pixel_values = []
        batch_pixel_video_values = []
        batch_video_thw = []
        batch_image_thw = []
        batch_second_per_grid_ts = []
        
        batch_image_files = []
        
        for example in examples:
            keys = example.keys()
            if "pixel_values_videos" in keys:
                batch_pixel_video_values.append(example["pixel_values_videos"])
                batch_video_thw.append(example["video_grid_thw"])
            elif "pixel_values" in keys:
                batch_pixel_values.append(example["pixel_values"])
                batch_image_thw.append(example["image_grid_thw"])
            
            if "image_files" in keys:
                batch_image_files.append(example["image_files"])
            
            batch_input_ids.append(example["input_ids"])
            batch_label_ids.append(example["labels"])

            if "second_per_grid_ts" in keys:
                batch_second_per_grid_ts.extend(example["second_per_grid_ts"])
        
        input_ids = pad_sequence(
            batch_input_ids, padding_side='right', padding_value=self.pad_token_id
        )

        attention_mask = input_ids != self.pad_token_id
        labels = pad_sequence(batch_label_ids, padding_side='right', padding_value=IGNORE_INDEX)

        data_dict = {
            'input_ids': input_ids,
            'labels': labels,
            'attention_mask': attention_mask,
        }

        if len(batch_pixel_values) > 0:
            pixel_values = torch.cat(batch_pixel_values, dim=0)
            image_thw = torch.cat(batch_image_thw, dim=0)
            data_dict["pixel_values"] = pixel_values
            data_dict["image_grid_thw"] = image_thw

        if len(batch_pixel_video_values) > 0:
            pixel_video_values = torch.cat(batch_pixel_video_values, dim=0)
            video_thw = torch.cat(batch_video_thw, dim=0)
            data_dict["pixel_values_videos"] = pixel_video_values
            data_dict["video_grid_thw"] = video_thw

        if len(batch_second_per_grid_ts) > 0:
            data_dict["second_per_grid_ts"] = batch_second_per_grid_ts
            
        if len(batch_image_files) > 0:
            data_dict["image_files"] = batch_image_files

        return data_dict

def replace_image_tokens(input_string, is_video=False):
    if is_video:
        pattern = r'\n?' + re.escape(LLAVA_VIDEO_TOKEN) + r'\n?'
        replacement = VISION_START_TOKEN + DEFAULT_VIDEO_TOKEN + VISION_END_TOKEN
    else:
        pattern = r'\n?' + re.escape(LLAVA_IMAGE_TOKEN) + r'\n?'
        replacement = VISION_START_TOKEN + DEFAULT_IMAGE_TOKEN + VISION_END_TOKEN

    return re.sub(pattern, replacement, input_string)

def llava_to_openai(conversations, is_video=False):
    role_mapping = {"human": "user", "gpt": "assistant"}

    transformed_data = []
    for conversation in conversations:
        transformed_content = replace_image_tokens(conversation["value"], is_video=is_video)
        transformed_entry = {
            "role": role_mapping.get(conversation["from"], conversation["from"]),
            "content": transformed_content,
        }
        transformed_data.append(transformed_entry)

    return transformed_data

def make_supervised_data_module(model_id, processor, data_args, anchor_model_id):
    """Make dataset and collator for supervised fine-tuning."""
    sft_dataset = SupervisedDataset(
        data_path=data_args.data_path, processor=processor, data_args=data_args, model_id=model_id, anchor_model_id=anchor_model_id
    )
    data_collator = DataCollatorForSupervisedDataset(pad_token_id=processor.tokenizer.pad_token_id)

    return dict(train_dataset=sft_dataset,
                eval_dataset=None,
                data_collator=data_collator)