src/gp_data.py

import os
import re
import ast
import math
import yaml
import warnings
from datetime import datetime
from dataclasses import dataclass, field
from collections import defaultdict
from typing import Any, Callable, Optional, Union, Sized, Dict, Tuple, List, Literal, Type

import numpy as np
import torch
from torch import nn
import torch.nn.functional as F

import datasets

from PIL import Image

from trl import ModelConfig, ScriptArguments, TrlParser, get_peft_config
from trl.models import unwrap_model_for_generation

from transformers import (
    TrainingArguments, 
    Trainer,
    GenerationConfig,
)
from transformers.modeling_utils import PreTrainedModel
from transformers.utils import (
    is_safetensors_available, 
    is_peft_available
)

if is_safetensors_available():
    import safetensors.torch
from peft import PeftConfig, get_peft_model, PeftModel
from accelerate.utils import is_peft_model, set_seed

from qwen_vl_utils import process_vision_info

from src.model.vlm_backbone.qwen2_5_vl_gp.process_gp import Qwen2_5_VL_GP_Processor

from transformers.trainer import (
    logger,
    TRAINING_ARGS_NAME,
    CONFIG_NAME,
    ADAPTER_WEIGHTS_NAME,
    ADAPTER_SAFE_WEIGHTS_NAME,
    WEIGHTS_NAME,
    WEIGHTS_INDEX_NAME,
    SAFE_WEIGHTS_NAME,
    SAFE_WEIGHTS_INDEX_NAME,
    FSDP_MODEL_NAME,
)

from src.model.vlm_backbone.qwen2_5_vl_gp.warppers import debug_calls
from src.utils_gp import (
    LLMClient, 
    norm_bboxes, 
    extract_one_bbox_from_str, 
    cal_paired_ious,
    print_rank0
)


# ---------- Datasets ----------

QUERY_KEY = "query"
IMG_PATH_KEY = "img_path"
ANSWER_KEY = "answer"
NORMED_BBOXES_KEY = "normed_bboxes"
SCORE_FUNCS_KEY = "score_funcs"

REMAIN_KEYS = [
    QUERY_KEY,
    IMG_PATH_KEY,
    NORMED_BBOXES_KEY,
    ANSWER_KEY,
    SCORE_FUNCS_KEY,
]

MAPPER_REGISTRY: Dict[str, Callable] = {}
FILTER_REGISTRY: Dict[str, Callable] = {}

def register_mappers():
    def wrapper(func):
        name = func.__name__.replace("_dataset_mapper", "")
        MAPPER_REGISTRY[name] = func
        return func
    return wrapper

def register_filters():
    def wrapper(func):
        name = func.__name__.replace("_dataset_filter", "")
        FILTER_REGISTRY[name] = func
        return func
    return wrapper


@register_mappers()
def cot_train_dataset_mapper(one_data, **kwargs):
    query = one_data['question']
    if 'prompt' in kwargs:
        query = kwargs['prompt'].format(query)
    answer = one_data['answer']
    image = one_data['image']
    dataset = one_data['dataset']
    img_path = os.path.join(kwargs['img_dir'], "cot", dataset, image)
    bboxes = one_data['bboxs']
    return {
        QUERY_KEY: query,
        ANSWER_KEY: answer,
        IMG_PATH_KEY: img_path,
        NORMED_BBOXES_KEY: bboxes,
        SCORE_FUNCS_KEY: kwargs['score_funcs']
    }    
    

@register_mappers()
def cot_train_fullmask_dataset_mapper(one_data, **kwargs):
    query = one_data['question']
    if 'prompt' in kwargs:
        query = kwargs['prompt'].format(query)
    answer = one_data['answer']
    image = one_data['image']
    dataset = one_data['dataset']
    img_path = os.path.join(kwargs['img_dir'], "cot", dataset, image)
    normed_bboxes = [[0.0, 0.0, 1.0, 1.0]]
    return {
        QUERY_KEY: query,
        ANSWER_KEY: answer,
        IMG_PATH_KEY: img_path,
        NORMED_BBOXES_KEY: normed_bboxes,
        SCORE_FUNCS_KEY: kwargs['score_funcs']
    }    
    
    
@register_mappers()
def norm_bboxes_dataset_mapper(one_data, **kwargs):
    bboxes = one_data.pop(NORMED_BBOXES_KEY)
    if 'width' in one_data:
        width = one_data['width']
        height = one_data['height']
    else:
        img_path = one_data[IMG_PATH_KEY]
        img_pil = Image.open(img_path)
        width, height = img_pil.size
        img_pil.close()
    normed_bboxes = norm_bboxes(bboxes, height, width, bbox_type=kwargs['bbox_type'])
    one_data[NORMED_BBOXES_KEY] = normed_bboxes
    return one_data

    
@register_filters()
def image_exist_dataset_filter(one_data, **kwargs):
    img_path = one_data[IMG_PATH_KEY]
    try:
        img = Image.open(img_path)
        img.close()
        return True
    except (FileNotFoundError, OSError) as e:
        print_rank0(f"Image not found or invalid: {img_path}. Error: {e}")
        return False
    except Exception as e:
        print_rank0(f"Unexpected error while checking image: {img_path}. Error: {e}")
        return False
    
@register_filters()
def inputs_seq_length_dataset_filter(one_data, **kwargs):
    processor = kwargs['processor']
    max_input_seq_length = kwargs.get('max_input_seq_length', None)
    max_input_remain_seq_length = kwargs.get('max_input_remain_seq_length', None)
    if max_input_seq_length is None and max_input_remain_seq_length is None:
        return True
    img_path = one_data[IMG_PATH_KEY]
    query = one_data[QUERY_KEY]
    normed_bboxes = [one_data[NORMED_BBOXES_KEY]] if max_input_remain_seq_length is not None else None
    messages = [[{"role": "user", "content": [{"type": "image", "image": img_path}, {"type": "text", "text": query}]}]]
    text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    image_inputs, video_inputs = process_vision_info(messages)
    inputs = processor(
            text=text,
            images=image_inputs,
            videos=video_inputs,
            normed_bboxes=normed_bboxes,
            padding=True,
            return_tensors="pt",
        )
    seq_length = inputs.input_ids.shape[1]
    if max_input_seq_length is not None and seq_length > max_input_seq_length:
        return False
    
    if max_input_remain_seq_length is not None:
        ref_token_masks = inputs.ref_token_masks[0]
        reduced_num = ref_token_masks.numel() - ref_token_masks.sum().item()
        remain_seq_length = seq_length - reduced_num
        if remain_seq_length > max_input_remain_seq_length:
            return False
    return True


# ---------- Loss ----------

LOSS_REGISTRY: Dict[str, Type[nn.Module]] = {}

def register_loss(loss_class):
    name = loss_class.__name__
    if name in LOSS_REGISTRY:
        raise ValueError(f"Loss class '{name}' is already registered.")
    LOSS_REGISTRY[name] = loss_class
    return loss_class


@register_loss
class DiceLoss(nn.Module):
    def __init__(self, epsilon: float = 1e-6, **kwargs):
        super().__init__()
        self.epsilon = epsilon

    def forward(self, 
                image_token_mask_logits: List[torch.Tensor], 
                ref_token_masks: List[torch.Tensor]
               ) -> torch.Tensor:
        if not isinstance(image_token_mask_logits, list) or not isinstance(ref_token_masks, list):
            raise TypeError("Inputs must be lists of tensors.")
        if len(image_token_mask_logits) != len(ref_token_masks):
            raise ValueError(f"Input lists must have the same length, but got "
                             f"{len(image_token_mask_logits)} and {len(ref_token_masks)}")
        if len(image_token_mask_logits) == 0:
            return torch.tensor(0.0, device=image_token_mask_logits[0].device if image_token_mask_logits else None) 

        batch_size = len(image_token_mask_logits)
        total_dice_loss = 0.0

        for i in range(batch_size):
            pred_mask_1d = image_token_mask_logits[i].flatten().sigmoid()
            gt_mask_1d = ref_token_masks[i].flatten().to(pred_mask_1d.device, dtype=torch.float)
            intersection = (pred_mask_1d * gt_mask_1d).sum()
            pred_sum = pred_mask_1d.sum()
            gt_sum = gt_mask_1d.sum()
            dice_coefficient = (2.0 * intersection + self.epsilon) / (pred_sum + gt_sum + self.epsilon)
            total_dice_loss += (1.0 - dice_coefficient)

        return total_dice_loss / batch_size


@register_loss
class BCELoss(nn.Module):
    def ___init__(self, **kwargs):
        super(BCELoss, self).__init__()
        
    def forward(self, 
                image_token_mask_logits: List[torch.Tensor], 
                ref_token_masks: List[torch.Tensor]
               ) -> torch.Tensor:
        
        batch_size = len(image_token_mask_logits)
        total_bce_loss = 0.0
        for i in range(batch_size):
            pred_mask_1d = image_token_mask_logits[i].flatten()
            gt_mask_1d = ref_token_masks[i].flatten().to(pred_mask_1d.device)
            bce_loss = F.binary_cross_entropy_with_logits(
                pred_mask_1d.float(),
                gt_mask_1d.float(),
            )
            total_bce_loss += bce_loss
        return total_bce_loss / batch_size


@register_loss
class MaskLoss(nn.Module):
    def __init__(self, 
                 dice_weight: float = 0.5,
                 bce_weight: float = 0.5,
                 epsilon: float = 1e-6,
                 **kwargs):
        super().__init__()
        self.dice_loss = DiceLoss(epsilon=epsilon)
        self.bce_loss = BCELoss()
        self.dice_weight = dice_weight
        self.bce_weight = bce_weight
    
    def forward(self, image_token_mask_logits: List[torch.Tensor],
                ref_token_masks: List[torch.Tensor]
               ) -> torch.Tensor:
        dice_loss = self.dice_loss(image_token_mask_logits, ref_token_masks)
        bce_loss = self.bce_loss(image_token_mask_logits, ref_token_masks)
        return self.dice_weight * dice_loss + self.bce_weight * bce_loss


# ---------- (Stub) Score functions (for YAML compatibility) ----------

SCORE_REGISTRY: Dict[str, Callable] = {}

def register_score():
    def wrapper(func):
        name = func.__name__.replace("_score", "")
        SCORE_REGISTRY[name] = func
        return func
    return wrapper

@register_score()
def llm_score(query, completion, answer, args):
    """
    YAML 里可能写了 'score_funcs: [llm]'。本工程不使用这些分数，返回 0 占位即可。
    """
    # 返回与 batch 大小一致的 0 分
    if isinstance(query, list):
        return [0.0] * len(query)
    return [0.0]


# ---------- Dataset & Collator & Sampler ----------

def _resolve_rel_path(rel_path: str, base_dir: str) -> str:
    """
    Resolve a relative path against base_dir; if not found, try parent dirs up to 4 levels.
    """
    if os.path.isabs(rel_path):
        return rel_path
    candidates = [os.path.join(base_dir, rel_path)]
    parent = base_dir
    for _ in range(4):
        parent = os.path.dirname(parent)
        if not parent or parent in ("/", ""):
            break
        candidates.append(os.path.join(parent, rel_path))
    for cand in candidates:
        if os.path.exists(cand):
            return cand
    return candidates[0]


class GPDataset(torch.utils.data.Dataset):
    """
    A PyTorch Dataset that loads and combines multiple datasets
    based on a YAML configuration file. It handles sampling
    and applies specified mapping functions.
    """
    @classmethod
    def _load_config(cls, config_path: str) -> Dict[str, Any]:
        print_rank0(f"Loading configuration from: {config_path}")
        try:
            with open(config_path, 'r', encoding='utf-8') as f:
                conf = yaml.safe_load(f)
            if conf is None:
                raise ValueError("YAML config is empty.")

            base_dir = os.path.dirname(config_path)
            # 允许传“顶层训练配置”：里面用 train_dataset 指向真正的数据清单
            if 'datasets' not in conf:
                if 'train_dataset' in conf:
                    ds_yaml = _resolve_rel_path(conf['train_dataset'], base_dir)
                    print_rank0(f"Loading dataset config from: {ds_yaml}")
                    with open(ds_yaml, 'r', encoding='utf-8') as f:
                        conf2 = yaml.safe_load(f)
                    if conf2 is None or 'datasets' not in conf2:
                        raise ValueError(f"'{ds_yaml}' missing 'datasets' key.")
                    conf = conf2
                    base_dir = os.path.dirname(ds_yaml)
                else:
                    raise ValueError("YAML config is missing both 'datasets' and 'train_dataset' keys.")

            conf['__root_dir__'] = base_dir
            print_rank0("Configuration loaded successfully.")
            return conf
        except Exception as e:
            print_rank0(f"Failed to load config: {e}")
            raise

    @classmethod
    def _apply_sampling(cls, dataset: datasets.Dataset, strategy: Optional[str], seed: Optional[int] = None) -> datasets.Dataset:
        """Applies sampling strategy to a dataset."""
        if not strategy:
            print_rank0("No sampling strategy specified, using full dataset.")
            return dataset

        try:
            parts = strategy.split(':')
            if len(parts) != 2:
                raise ValueError(f"Invalid sampling strategy format: '{strategy}'. Expected 'type:value'.")
            strat_type, strat_value = parts[0].lower(), parts[1]
            num_samples = int(strat_value)
            total_size = len(dataset)
            if num_samples <= 0:
                raise ValueError(f"Sampling value must be positive, got: {num_samples} [{strategy}]")
            num_samples = min(num_samples, total_size)

            print_rank0(f"Applying sampling: {strategy} ({num_samples} samples) to dataset of size {total_size}")

            if strat_type == "first":
                return dataset.select(range(num_samples))
            elif strat_type == "end":
                start_index = max(0, total_size - num_samples)
                return dataset.select(range(start_index, total_size))
            elif strat_type == "random":
                shuffled_dataset = dataset.shuffle(seed=seed)
                return shuffled_dataset.select(range(num_samples))
            else:
                print_rank0(f"Warning: Unknown sampling strategy type: '{strat_type}'. Using full dataset.")
                return dataset
        except ValueError as e:
            print_rank0(f"Error parsing sampling strategy '{strategy}': {e}. Using full dataset.")
            return dataset
        except Exception as e:
            print_rank0(f"An unexpected error occurred during sampling: {e}. Using full dataset.")
            return dataset
        
    @classmethod
    def _all_processed_datasets(cls, config, processor, args):
        root_dir = config.get('__root_dir__', os.getcwd())
        all_processed_datasets: Dict[str, datasets.Dataset] = {}
        for i, dataset_config in enumerate(config['datasets']):
            print_rank0(f"\nProcessing dataset entry {i+1}/{len(config['datasets'])}...")
            json_path = dataset_config.get('json_path')
            if not json_path:
                print_rank0(f"Warning: Skipping dataset entry {i+1} due to missing 'json_path'.")
                continue
            json_path = _resolve_rel_path(json_path, root_dir)

            base_name = '.'.join(os.path.basename(json_path).split('.')[:-1])
            dataset_name = dataset_config.get('dataset_name', base_name)

            sampling_strategy = dataset_config.get('sampling_strategy', None)
            sampling_seed = dataset_config['sampling_seed'] if 'sampling_seed' in dataset_config else getattr(args, 'sampling_seed', 42)

            mapper_name = dataset_config.get('mapper')
            bbox_type = dataset_config.get('bbox_type')

            # img_dir: 优先用数据 YAML 里的；否则尝试 args.img_dir（可能不存在）
            if 'img_dir' in dataset_config:
                img_dir = _resolve_rel_path(dataset_config['img_dir'], root_dir)
            else:
                img_dir = getattr(args, 'img_dir', None)
                if img_dir is not None:
                    img_dir = _resolve_rel_path(img_dir, root_dir)

            additional_mappers = dataset_config.get('additional_mappers', [])
            score_funcs = dataset_config.get('score_funcs', [])
            prompt = dataset_config.get('prompt', None)

            max_input_seq_length = dataset_config['max_input_seq_length'] if 'max_input_seq_length' in dataset_config else getattr(args, 'max_input_seq_length', None)
            max_input_remain_seq_length = dataset_config['max_input_remain_seq_length'] if 'max_input_remain_seq_length' in dataset_config else getattr(args, 'max_input_remain_seq_length', None)
            
            # 安全处理 score_funcs：过滤未注册的（不报错，只警告）
            if score_funcs:
                filtered = []
                for sf in score_funcs:
                    if sf in SCORE_REGISTRY:
                        filtered.append(sf)
                    else:
                        print_rank0(f"Warning: Score function '{sf}' not registered. Will ignore.")
                score_funcs = filtered

            try:
                print_rank0(f"Loading raw data from: {json_path}")
                raw_dataset = datasets.load_dataset('json', data_files=json_path, split='train')
                print_rank0(f"Loaded {len(raw_dataset)} examples raw.")

                sampled_dataset = cls._apply_sampling(raw_dataset, sampling_strategy, sampling_seed)
                if len(sampled_dataset) == 0:
                    print_rank0("Dataset is empty after sampling, skipping.")
                    continue
                print_rank0(f"Dataset size after sampling: {len(sampled_dataset)}")

                mapper_func = MAPPER_REGISTRY[mapper_name]
                print_rank0(f"Applying mapper: '{mapper_name}'")
                mapper_kwargs = {
                    'img_dir': img_dir,
                    'score_funcs': score_funcs,
                }
                if prompt is not None:
                    mapper_kwargs['prompt'] = prompt
                print_rank0(f"Mapper arguments: {mapper_kwargs}")

                processed_dataset = sampled_dataset.map(
                    mapper_func,
                    num_proc=8,
                    fn_kwargs=mapper_kwargs,
                )

                processed_dataset = processed_dataset.remove_columns(
                    [col for col in processed_dataset.column_names if col not in REMAIN_KEYS]
                )
                    
                print_rank0("Applying dataset filter: 'image_exist_dataset_filter'")
                processed_dataset = processed_dataset.filter(
                    image_exist_dataset_filter,
                    num_proc=8,
                    fn_kwargs={}
                )
                print_rank0(f"Processed dataset size after image_exist_dataset_filter: {len(processed_dataset)}")
                
                if max_input_seq_length is not None or max_input_remain_seq_length is not None:
                    processed_dataset = processed_dataset.filter(
                        inputs_seq_length_dataset_filter,
                        num_proc=8,
                        fn_kwargs={
                            'processor': processor,
                            'max_input_seq_length': max_input_seq_length,
                            'max_input_remain_seq_length': max_input_remain_seq_length,
                        }
                    )
                    print_rank0(f"Processed dataset size after inputs_seq_length_dataset_filter: {len(processed_dataset)}")
                
                for additional_mapper in additional_mappers:
                    mapper_func = MAPPER_REGISTRY[additional_mapper]
                    print_rank0(f"Applying additional mapper: '{additional_mapper}'")
                    processed_dataset = processed_dataset.map(
                        mapper_func,
                        num_proc=8,
                        fn_kwargs={
                            'bbox_type': bbox_type,
                        }
                    )
                print_rank0(f"Processed dataset size: {len(processed_dataset)}")
                if len(processed_dataset) == 0:
                    print_rank0(f"Warning: Processed dataset {dataset_name} is empty after mapping. Skipping.")
                    continue

                if dataset_name in all_processed_datasets:
                    dataset_name_with_uuid = f"{dataset_name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
                    print_rank0(f"Warning: Dataset name '{dataset_name}' already exists. Renaming to '{dataset_name_with_uuid}'")
                    all_processed_datasets[dataset_name_with_uuid] = processed_dataset
                else:
                    all_processed_datasets[dataset_name] = processed_dataset                

            except FileNotFoundError:
                print_rank0(f"Error: Data file not found for dataset entry {i+1}: {json_path}. Skipping.")
            except Exception as e:
                print_rank0(f"Error processing dataset entry {i+1} ({json_path}): {e}. Skipping.")
                
        return all_processed_datasets
        

    def __init__(self, config_path: str, processor: Qwen2_5_VL_GP_Processor, script_args: Optional[Any] = None):
        """
        Initializes the GPDataset.

        Args:
            config_path (str): Path to the YAML configuration file.
            processor (Qwen2_5_VL_GP_Processor): Processor for handling text and vision data.
            script_args (Any, optional): Additional arguments passed from the script
                                         (e.g., training args, could contain seed). Defaults to None.
        """
        super().__init__()
        self.args = script_args
        self.config = self._load_config(config_path)
        self.processor = processor
        all_processed_datasets = self._all_processed_datasets(self.config, self.processor, self.args)
        if all_processed_datasets:
            print_rank0(f"\nConcatenating {len(all_processed_datasets)} processed dataset(s)...")
            self.final_dataset = datasets.concatenate_datasets(list(all_processed_datasets.values()))
            if len(self.final_dataset) == 0:
                raise ValueError("Final dataset is empty after concatenation.")
            print_rank0(f"Final combined dataset size: {len(self.final_dataset)}")
            print_rank0(f"Final dataset features: {self.final_dataset.features}")
        else:
            raise ValueError("No datasets were successfully processed. Please check your configuration.")
            self.final_dataset = None

    def __len__(self) -> int:
        return len(self.final_dataset) if self.final_dataset else 0

    def __getitem__(self, index: int) -> Dict[str, Any]:
        if self.final_dataset is None:
            raise IndexError("Dataset is not initialized or is empty.")
        if not 0 <= index < len(self.final_dataset):
             raise IndexError(f"Index {index} out of bounds for dataset of size {len(self.final_dataset)}")
        return self.final_dataset[index]
    
    
    @classmethod
    def get_processed_dataset_dict(cls, config_path: str, processor: Qwen2_5_VL_GP_Processor, script_args: Optional[Any] = None) -> Dict[str, datasets.Dataset]:
        config = cls._load_config(config_path)
        all_processed_datasets = cls._all_processed_datasets(config, processor, script_args)
        return all_processed_datasets



class GPCollator:
    def __init__(self, processor, is_sft):
        self.processor = processor
        self.is_sft = is_sft
        self.im_start_id = self.processor.tokenizer.encode("<|im_start|>")[0]
        
    def _prepare_labels_from_input_ids(self, input_ids):
        B, L = input_ids.shape
        labels = input_ids.clone()
        mask = input_ids == self.im_start_id
        flipped_mask = mask.flip(dims=(1,))
        first_idx_in_flipped = torch.argmax(flipped_mask.int(), dim=1)
        last_pos = (L - 1) - first_idx_in_flipped
        mask_until_idx = last_pos + 3
        mask_until_idx = torch.clamp(mask_until_idx, max=L)
        arange_l = torch.arange(L, device=input_ids.device).expand(B, -1)
        modification_mask = arange_l < mask_until_idx.unsqueeze(1)
        labels[modification_mask] = -100
        return labels
        
    def __call__(self, features):
        messages = []
        normed_bboxes = []
        answers = []
        querys = []
        score_funcs = []
        for feature in features:
            query = feature[QUERY_KEY]
            answer = feature[ANSWER_KEY]
            img_path = feature[IMG_PATH_KEY]
            if self.is_sft:
                messages.append([{"role": "user", "content": [{"type": "image", "image": img_path}, {"type": "text", "text": query}]}, {"role": "assistant", "content": [{"type": "text", "text": answer}]}])
            else:
                messages.append([{"role": "user", "content": [{"type": "image", "image": img_path}, {"type": "text", "text": query}]}])
            normed_bboxes.append(feature[NORMED_BBOXES_KEY])
            querys.append(query)
            answers.append(answer)
            score_funcs.append(feature[SCORE_FUNCS_KEY])
        
        text = self.processor.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=(not self.is_sft)
        )
        image_inputs, video_inputs = process_vision_info(messages)
        inputs = self.processor(
            text=text,
            normed_bboxes=normed_bboxes,
            images=image_inputs,
            videos=video_inputs,
            padding=True,
            return_tensors="pt",
        )
        
        if self.is_sft:
            labels = self._prepare_labels_from_input_ids(inputs.input_ids)
            inputs["labels"] = labels
        
        inputs[QUERY_KEY] = querys
        inputs[ANSWER_KEY] = answers
        inputs[SCORE_FUNCS_KEY] = score_funcs
        return inputs