training/pretrain/pretrain_dataset_adaptive.py

#!/usr/bin/env python3
"""
自适应Prompt的数据集加载器
使用数据中的user_prompt字段，而不是固定的prompt模板
"""

import pickle
from pathlib import Path
from typing import Dict, List, Optional, Tuple
import torch
from torch.utils.data import Dataset
from PIL import Image
import random


class AdaptivePretrainDataset(Dataset):
    """
    自适应Prompt预训练数据集
    每个样本都有自己的user_prompt，根据annotation长度定制
    
    Args:
        data_file: pretrain_data_adaptive.pkl路径
        split: 'train', 'val', 或 'test'
        tasks: 任务列表
        curriculum_stage: 0=easy, 1=medium, 2=hard, 3=all
        use_system_prompt: 是否使用system prompt
    """
    
    # System prompts (任务级别)
    SYSTEM_PROMPTS = {
        "scene_understanding": "You are an expert driving scene analyzer. Describe the environment accurately.",
        "binary_detection": "You are a traffic safety AI. Detect abnormal driving situations.",
        "accident_description": "You are an accident analysis AI. Answer based on the question asked.",  # 更通用
        "sequence_prediction": "You are a temporal driving AI. Analyze video sequences for accident prediction."
    }
    
    def __init__(
        self,
        data_file: str,
        split: str = "train",
        tasks: List[str] = None,
        curriculum_stage: int = 3,
        use_system_prompt: bool = True
    ):
        self.split = split
        self.tasks = tasks or ["task1", "task2", "task3", "task4"]
        self.curriculum_stage = curriculum_stage
        self.use_system_prompt = use_system_prompt
        
        # 加载数据
        with open(data_file, "rb") as f:
            all_data = pickle.load(f)
        
        split_data = all_data[split]
        
        # 收集样本
        self.samples = []
        
        task_map = {
            "task1": "task1_scene_understanding",
            "task2": "task2_binary_detection",
            "task3": "task3_accident_description",
            "task4": "task4_sequence_prediction"
        }
        
        for task in self.tasks:
            if task in task_map:
                task_samples = split_data.get(task_map[task], [])
                
                # Curriculum filtering
                if curriculum_stage < 3:
                    difficulty_map = {0: "easy", 1: "medium", 2: "hard"}
                    target_difficulty = difficulty_map[curriculum_stage]
                    task_samples = [
                        s for s in task_samples 
                        if s.get("difficulty", "easy") == target_difficulty
                    ]
                
                self.samples.extend(task_samples)
        
        # Shuffle
        if split == "train":
            random.shuffle(self.samples)
        
        print(f"{'='*70}")
        print(f"数据集加载: {split}")
        print(f"Curriculum Stage: {curriculum_stage} ({['easy', 'medium', 'hard', 'all'][curriculum_stage]})")
        print(f"任务: {tasks}")
        print(f"样本数: {len(self.samples)}")
        
        # 统计
        from collections import Counter
        task_dist = Counter(s["task"] for s in self.samples)
        print(f"\n任务分布:")
        for task, count in task_dist.items():
            print(f"  {task}: {count}")
        
        # 统计短/长标注
        if curriculum_stage == 3 and ("task3" in self.tasks or "task4" in self.tasks):
            short_count = sum(
                1 for s in self.samples 
                if s["task"] in ["accident_description", "sequence_prediction"]
                and s["metadata"].get("is_short_annotation", False)
            )
            detailed_count = sum(
                1 for s in self.samples 
                if s["task"] in ["accident_description", "sequence_prediction"]
                and not s["metadata"].get("is_short_annotation", False)
            )
            
            if short_count + detailed_count > 0:
                print(f"\nAnnotation分布 (任务3&4):")
                print(f"  短标注 (<20字符): {short_count}")
                print(f"  详细标注 (>=20字符): {detailed_count}")
        
        # 难度分布
        if curriculum_stage == 3:
            diff_dist = Counter(s.get("difficulty", "unknown") for s in self.samples)
            print(f"\n难度分布:")
            for diff, count in diff_dist.items():
                print(f"  {diff}: {count}")
        
        print("=" * 70)
    
    def __len__(self):
        return len(self.samples)
    
    def __getitem__(self, idx):
        sample = self.samples[idx]
        task_type = sample["task"]
        
        # 获取system prompt (任务级别)
        system_prompt = self.SYSTEM_PROMPTS[task_type] if self.use_system_prompt else ""
        
        # 使用样本中的user_prompt (自适应)
        user_prompt = sample.get("user_prompt", "")
        
        if task_type in ["scene_understanding", "binary_detection", "accident_description"]:
            # 单帧任务
            image = Image.open(sample["image_path"]).convert("RGB")
            
            return {
                "task": task_type,
                "subtask": sample.get("subtask", task_type),
                "image": image,
                "system_prompt": system_prompt,
                "user_prompt": user_prompt,  # 自适应prompt
                "label": sample["label"],
                "difficulty": sample.get("difficulty", "unknown"),
                "metadata": sample["metadata"]
            }
        
        elif task_type == "sequence_prediction":
            # 序列任务
            images = []
            for img_path in sample["image_sequence"]:
                img = Image.open(img_path).convert("RGB")
                images.append(img)
            
            return {
                "task": task_type,
                "subtask": sample.get("subtask", task_type),
                "image_sequence": images,
                "system_prompt": system_prompt,
                "user_prompt": user_prompt,  # 自适应prompt
                "label": sample["label"],
                "difficulty": sample.get("difficulty", "unknown"),
                "metadata": sample["metadata"]
            }
        
        else:
            raise ValueError(f"未知任务类型: {task_type}")


def collate_fn_adaptive(batch):
    """
    自适应collate函数
    每个样本有自己的user_prompt
    """
    single_frame_batch = []
    sequence_batch = []
    
    for item in batch:
        if item["task"] in ["scene_understanding", "binary_detection", "accident_description"]:
            single_frame_batch.append(item)
        elif item["task"] == "sequence_prediction":
            sequence_batch.append(item)
    
    result = {}
    
    # 单帧任务
    if single_frame_batch:
        result["single_frame"] = {
            "task": [x["task"] for x in single_frame_batch],
            "subtask": [x["subtask"] for x in single_frame_batch],
            "images": [x["image"] for x in single_frame_batch],
            "system_prompts": [x["system_prompt"] for x in single_frame_batch],
            "user_prompts": [x["user_prompt"] for x in single_frame_batch],  # 每个样本不同
            "labels": [x["label"] for x in single_frame_batch],
            "difficulties": [x["difficulty"] for x in single_frame_batch],
            "metadata": [x["metadata"] for x in single_frame_batch]
        }
    
    # 序列任务
    if sequence_batch:
        result["sequence"] = {
            "task": [x["task"] for x in sequence_batch],
            "subtask": [x["subtask"] for x in sequence_batch],
            "image_sequences": [x["image_sequence"] for x in sequence_batch],
            "system_prompts": [x["system_prompt"] for x in sequence_batch],
            "user_prompts": [x["user_prompt"] for x in sequence_batch],  # 每个样本不同
            "labels": [x["label"] for x in sequence_batch],
            "difficulties": [x["difficulty"] for x in sequence_batch],
            "metadata": [x["metadata"] for x in sequence_batch]
        }
    
    return result


# ============ 测试代码 ============
if __name__ == "__main__":
    from torch.utils.data import DataLoader
    
    data_file = "PROJECT_ROOT/data/dataset/pretrain/train/pretrain_data_adaptive.pkl"
    
    print("\n" + "=" * 70)
    print("测试自适应Prompt数据集")
    print("=" * 70)
    
    # 创建数据集
    dataset = AdaptivePretrainDataset(
        data_file=data_file,
        split="train",
        tasks=["task1", "task2", "task3", "task4"],
        curriculum_stage=3
    )
    
    loader = DataLoader(
        dataset,
        batch_size=4,
        shuffle=False,
        num_workers=0,
        collate_fn=collate_fn_adaptive
    )
    
    # 测试一个batch
    batch = next(iter(loader))
    
    print("\n" + "=" * 70)
    print("Batch示例")
    print("=" * 70)
    
    if "single_frame" in batch:
        sf = batch["single_frame"]
        print(f"\n单帧任务: {len(sf['images'])} 样本")
        
        for i in range(len(sf['task'])):
            print(f"\n样本 {i+1}:")
            print(f"  任务: {sf['task'][i]}")
            print(f"  难度: {sf['difficulties'][i]}")
            print(f"  System: {sf['system_prompts'][i][:60]}...")
            print(f"  User Prompt: {sf['user_prompts'][i]}")  # 注意每个都不同
            print(f"  Label: {sf['labels'][i][:60]}...")
            
            # 如果是事故描述任务，显示annotation长度
            if sf['task'][i] == 'accident_description':
                is_short = sf['metadata'][i].get('is_short_annotation', False)
                anno_len = sf['metadata'][i].get('annotation_length', 0)
                print(f"  Annotation: {'短' if is_short else '详细'} ({anno_len}字符)")
    
    if "sequence" in batch:
        seq = batch["sequence"]
        print(f"\n序列任务: {len(seq['image_sequences'])} 样本")
        
        for i in range(len(seq['task'])):
            print(f"\n样本 {i+1}:")
            print(f"  序列长度: {len(seq['image_sequences'][i])}")
            print(f"  难度: {seq['difficulties'][i]}")
            print(f"  User Prompt: {seq['user_prompts'][i]}")  # 注意每个都不同
            print(f"  Label: {seq['labels'][i][:60]}...")
            
            is_short = seq['metadata'][i].get('is_short_annotation', False)
            anno_len = seq['metadata'][i].get('annotation_length', 0)
            print(f"  Annotation: {'短' if is_short else '详细'} ({anno_len}字符)")
    
    print("\n✅ 数据集测试完成!")