data_loader.py

import torch
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader, IterableDataset
from datasets import load_dataset, concatenate_datasets, interleave_datasets
from typing import Dict, List, Optional, Any, Union
import random
import numpy as np
from tqdm import tqdm
import warnings
from PIL import Image
import requests
from io import BytesIO
from torchvision import transforms
import logging

# 设置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

warnings.filterwarnings("ignore", category=UserWarning)

from data_config import (
    PRETRAIN_DATASETS,
    POSTTRAIN_DATASETS,
    TEST_DATASETS,
    PRETRAIN_MIX,
    POSTTRAIN_MIX,
    PREPROCESSING_CONFIG,
    DATASET_CACHE_DIR,
    HF_CACHE_DIR
)

# 图像变换
image_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

class PreTrainDataset(IterableDataset):
    def __init__(
        self,
        mix_name: str = 'default',
        tokenizer=None,
        max_length: int = 2048,
        streaming: bool = True,
        seed: int = 42,
        max_samples: Optional[int] = None
    ):
        super().__init__()
        
        if tokenizer is None:
            raise ValueError("tokenizer cannot be None")
        
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.streaming = streaming
        self.seed = seed
        self.max_samples = max_samples
        self.samples_generated = 0
        
        # 获取混合配置
        if mix_name not in PRETRAIN_MIX:
            raise ValueError(f"Unknown mix: {mix_name}. Available: {list(PRETRAIN_MIX.keys())}")
        
        mix_config = PRETRAIN_MIX[mix_name]
        dataset_names = mix_config.get('datasets', [])
        weights = mix_config.get('weights', [])
        
        if not dataset_names:
            raise ValueError(f"No datasets found in mix: {mix_name}")
        
        logger.info(f"Loading pretrain mix: {mix_name}")
        logger.info(f"  Datasets: {dataset_names}")
        logger.info(f"  Weights: {weights}")
        
        # 加载数据集
        self.datasets = []
        self.probabilities = []
        
        for name, weight in zip(dataset_names, weights):
            if name not in PRETRAIN_DATASETS:
                logger.warning(f"Dataset {name} not found in PRETRAIN_DATASETS, skipping")
                continue
            
            config = PRETRAIN_DATASETS[name]
            try:
                ds = self._load_dataset(config)
                if ds is not None:
                    self.datasets.append((name, ds, config))
                    self.probabilities.append(weight)
                    logger.info(f"  Successfully loaded {name}")
            except Exception as e:
                logger.error(f"Error loading {name}: {e}")
                continue
        
        if not self.datasets:
            raise ValueError("No datasets loaded successfully")
        
        # 归一化概率
        total = sum(self.probabilities)
        self.probabilities = [p / total for p in self.probabilities]
        
        logger.info(f"Successfully loaded {len(self.datasets)} datasets")

    def _load_dataset(self, config: Dict):
        try:
            load_kwargs = {
                'path': config['hf_path'],
                'split': config.get('split', 'train'),
                'streaming': config.get('streaming', self.streaming),
                'cache_dir': HF_CACHE_DIR,
            }
            
            if 'config' in config:
                load_kwargs['name'] = config['config']
            
            ds = load_dataset(**load_kwargs)
            return ds
        except Exception as e:
            logger.error(f"Failed to load {config.get('hf_path', 'unknown')}: {e}")
            return None

    def _process_text_sample(self, sample: Dict, config: Dict) -> Optional[Dict]:
        try:
            text_field = config.get('text_field', 'text')
            text = sample.get(text_field, '')
            
            if not text or not isinstance(text, str):
                return None
            
            text = text.strip()
            if len(text) < 10:
                return None
            
            # Tokenize
            encoding = self.tokenizer(
                text,
                max_length=self.max_length,
                truncation=True,
                padding='max_length',
                return_tensors='pt'
            )
            
            return {
                'input_ids': encoding['input_ids'].squeeze(0),
                'attention_mask': encoding['attention_mask'].squeeze(0),
                'type': 'text'
            }
        except Exception as e:
            logger.debug(f"Error processing text sample: {e}")
            return None

    def _process_image_text_sample(self, sample: Dict, config: Dict) -> Optional[Dict]:
        try:
            text_field = config.get('text_field', 'caption')
            image_field = config.get('image_field', 'image')
            
            text = sample.get(text_field, '')
            image = sample.get(image_field)
            
            if not text or image is None:
                return None
            
            # 处理图像
            if isinstance(image, str):
                try:
                    response = requests.get(image, timeout=5)
                    image = Image.open(BytesIO(response.content)).convert('RGB')
                except Exception as img_error:
                    logger.debug(f"Failed to load image from URL: {img_error}")
                    return None
            elif isinstance(image, Image.Image):
                image = image.convert('RGB')
            else:
                return None
            
            # 转换图像
            image_tensor = image_transform(image)
            
            # Tokenize文本
            encoding = self.tokenizer(
                text,
                max_length=self.max_length,
                truncation=True,
                padding='max_length',
                return_tensors='pt'
            )
            
            return {
                'input_ids': encoding['input_ids'].squeeze(0),
                'attention_mask': encoding['attention_mask'].squeeze(0),
                'image': image_tensor,
                'type': 'image_text'
            }
        except Exception as e:
            logger.debug(f"Error processing image-text sample: {e}")
            return None

    def __iter__(self):
        """迭代器"""
        worker_info = torch.utils.data.get_worker_info()
        if worker_info is not None:
            # 多worker时设置不同的随机种子
            random.seed(self.seed + worker_info.id)
            np.random.seed(self.seed + worker_info.id)
        else:
            random.seed(self.seed)
            np.random.seed(self.seed)
        
        # 创建数据集迭代器
        iterators = [iter(ds) for _, ds, _ in self.datasets]
        self.samples_generated = 0
        
        while True:
            # 检查是否达到最大样本数
            if self.max_samples and self.samples_generated >= self.max_samples:
                break
            
            try:
                # 根据概率选择数据集
                idx = np.random.choice(len(self.datasets), p=self.probabilities)
                name, _, config = self.datasets[idx]
                
                # 从选中的数据集获取样本
                sample = next(iterators[idx])
                
                # 处理样本
                processed = None
                if config.get('type') in ['text', 'code']:
                    processed = self._process_text_sample(sample, config)
                elif config.get('type') == 'image_text':
                    processed = self._process_image_text_sample(sample, config)
                else:
                    logger.debug(f"Unknown type: {config.get('type')}")
                    continue
                
                if processed is not None:
                    self.samples_generated += 1
                    yield processed
                    
            except StopIteration:
                # 重新创建迭代器
                try:
                    iterators[idx] = iter(self.datasets[idx][1])
                except Exception as e:
                    logger.error(f"Failed to recreate iterator for dataset {idx}: {e}")
                    break
            except Exception as e:
                logger.debug(f"Error in iterator: {e}")
                continue


class PostTrainDataset(Dataset):
    def __init__(
        self,
        mix_name: str = 'default',
        tokenizer=None,
        max_length: int = 2048,
        max_samples: Optional[int] = None,
        split: str = 'train'
    ):
        super().__init__()
        
        if tokenizer is None:
            raise ValueError("tokenizer cannot be None")
        
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.split = split
        
        # 获取混合配置
        if mix_name not in POSTTRAIN_MIX:
            raise ValueError(f"Unknown mix: {mix_name}. Available: {list(POSTTRAIN_MIX.keys())}")
        
        mix_config = POSTTRAIN_MIX[mix_name]
        dataset_names = mix_config.get('datasets', [])
        weights = mix_config.get('weights', [])
        
        if not dataset_names:
            raise ValueError(f"No datasets found in mix: {mix_name}")
        
        logger.info(f"Loading posttrain mix: {mix_name}")
        logger.info(f"  Datasets: {dataset_names}")
        
        # 加载和合并数据集
        all_datasets = []
        
        for name in dataset_names:
            if name not in POSTTRAIN_DATASETS:
                logger.warning(f"Dataset {name} not found in POSTTRAIN_DATASETS")
                continue
            
            config = POSTTRAIN_DATASETS[name]
            try:
                load_kwargs = {
                    'path': config['hf_path'],
                    'split': split,
                    'streaming': config.get('streaming', False),
                    'cache_dir': HF_CACHE_DIR,
                }
                if 'data_files' in config:
                    load_kwargs['data_files'] = config['data_files']
                if 'config' in config:
                    load_kwargs['name'] = config['config']
                
                ds = load_dataset(**load_kwargs)
                
                # 限制样本数
                if config.get('max_samples'):
                    if hasattr(ds, 'take'):
                        ds = ds.take(config['max_samples'])
                    elif hasattr(ds, 'select'):
                        ds = ds.select(range(min(len(ds), config['max_samples'])))
                
                # 添加数据集标识
                def add_source(example):
                    example['_source'] = name
                    example['_config'] = config
                    return example
                
                ds = ds.map(add_source)
                all_datasets.append(ds)
                
                ds_len = len(ds) if hasattr(ds, '__len__') else 'streaming'
                logger.info(f"  Loaded {name}: {ds_len} samples")
                
            except Exception as e:
                logger.error(f"Error loading {name}: {e}")
                continue
        
        # 合并数据集
        if not all_datasets:
            raise ValueError("No datasets loaded successfully")
        
        if len(all_datasets) == 1:
            self.dataset = all_datasets[0]
        else:
            # 交织数据集
            probabilities = [w / sum(weights[:len(all_datasets)]) 
                           for w in weights[:len(all_datasets)]]
            self.dataset = interleave_datasets(
                all_datasets,
                probabilities=probabilities,
                seed=42,
                stopping_strategy='all_exhausted'
            )
        
        # 限制总样本数
        if max_samples and hasattr(self.dataset, '__len__'):
            actual_len = min(len(self.dataset), max_samples)
            self.dataset = self.dataset.select(range(actual_len))
        
        dataset_len = len(self.dataset) if hasattr(self.dataset, '__len__') else 'streaming'
        logger.info(f"Total samples: {dataset_len}")

    def _format_instruction(self, sample: Dict, config: Dict) -> str:
        """格式化instruction"""
        try:
            data_type = config.get('type', 'instruction')
            
            if data_type == 'instruction':
                instruction_field = config.get('instruction_field', 'instruction')
                input_field = config.get('input_field', 'input')
                context_field = config.get('context_field', None)
                
                instruction = sample.get(instruction_field, '')
                input_text = sample.get(input_field, '')
                context = sample.get(context_field, '') if context_field else ''
                
                # 构建prompt
                prompt_parts = [f"Instruction: {instruction}"]
                
                if context:
                    prompt_parts.append(f"Context: {context}")
                
                if input_text:
                    prompt_parts.append(f"Input: {input_text}")
                
                prompt_parts.append("Response:")
                return "\n".join(prompt_parts)
            
            elif data_type == 'conversation':
                if 'conversations' in sample:
                    conversations = sample['conversations']
                    if isinstance(conversations, list) and len(conversations) > 0:
                        dialogue = []
                        for conv in conversations[:-1]:
                            role = conv.get('from', 'user')
                            content = conv.get('value', '')
                            dialogue.append(f"{role}: {content}")
                        return "\n".join(dialogue) + "\nassistant:"
                
                elif 'messages' in sample:
                    # 标准消息格式
                    messages = sample['messages']
                    if isinstance(messages, list) and len(messages) > 0:
                        dialogue = []
                        for msg in messages[:-1]:
                            role = msg.get('role', 'user')
                            content = msg.get('content', '')
                            dialogue.append(f"{role}: {content}")
                        return "\n".join(dialogue) + "\nassistant:"
                
                # 如果没有标准格式，尝试使用text字段
                return sample.get('text', '')
            
            elif data_type == 'code_instruction':
                # 代码instruction格式
                instruction_field = config.get('instruction_field', 'instruction')
                instruction = sample.get(instruction_field, '')
                return f"### Instruction:\n{instruction}\n### Response:"
            
            elif data_type == 'multimodal_instruction':
                # 多模态instruction
                instruction_field = config.get('instruction_field', 'conversations')
                conversations = sample.get(instruction_field, [])
                if isinstance(conversations, list) and len(conversations) > 0:
                    # 提取对话历史（除了最后一条回复）
                    dialogue = []
                    for conv in conversations[:-1]:
                        role = conv.get('from', 'user')
                        content = conv.get('value', '')
                        dialogue.append(f"{role}: {content}")
                    return "\n".join(dialogue) + "\nassistant:"
                return ""
            
            else:
                return sample.get(config.get('instruction_field', 'text'), '')
        except Exception as e:
            logger.debug(f"Error formatting instruction: {e}")
            return ""

    def _get_response(self, sample: Dict, config: Dict) -> str:
        try:
            data_type = config.get('type', 'instruction')
            
            if data_type == 'instruction' or data_type == 'code_instruction':
                response_field = config.get('response_field', 'output')
                return sample.get(response_field, '')
            
            elif data_type == 'conversation':
                # 从对话中提取最后一条assistant的回复
                if 'conversations' in sample:
                    conversations = sample['conversations']
                    if isinstance(conversations, list) and len(conversations) > 0:
                        return conversations[-1].get('value', '')
                
                elif 'messages' in sample:
                    messages = sample['messages']
                    if isinstance(messages, list) and len(messages) > 0:
                        return messages[-1].get('content', '')
                
                return ""
            
            elif data_type == 'multimodal_instruction':
                instruction_field = config.get('instruction_field', 'conversations')
                conversations = sample.get(instruction_field, [])
                if isinstance(conversations, list) and len(conversations) > 0:
                    return conversations[-1].get('value', '')
                return ""
            
            else:
                response_field = config.get('response_field', 'output')
                return sample.get(response_field, '')
        except Exception as e:
            logger.debug(f"Error getting response: {e}")
            return ""

    def __len__(self):
        return len(self.dataset) if hasattr(self.dataset, '__len__') else 0

    def __getitem__(self, idx):
        try:
            sample = self.dataset[idx]
            
            # 获取配置
            if '_config' not in sample:
                logger.warning(f"Sample at index {idx} missing _config")
                return None
            
            config = sample['_config']
            
            # 格式化 instruction 和 response
            instruction_text = self._format_instruction(sample, config)
            response_text = self._get_response(sample, config)
            
            if not instruction_text or not response_text:
                return None
            
            pad_token_id = self.tokenizer.pad_token_id 
            if pad_token_id is None:
                pad_token_id = self.tokenizer.eos_token_id
            instruction_max_len = self.max_length // 2
            
            # Tokenize 不做 padding，手动处理
            instruction_enc = self.tokenizer(
                instruction_text,
                truncation=True,
                max_length=instruction_max_len,
                add_special_tokens=False, 
                return_tensors='pt'
            )
            instr_ids = instruction_enc['input_ids'].squeeze(0)
            
            # Instruction 手动 Padding
            instr_len = instr_ids.size(0)
            if instr_len < instruction_max_len:
                padding = torch.full((instruction_max_len - instr_len,), pad_token_id, dtype=torch.long)
                instr_ids = torch.cat([instr_ids, padding])
                
                instr_mask = torch.cat([torch.ones(instr_len, dtype=torch.long), torch.zeros(instruction_max_len - instr_len, dtype=torch.long)])
            else:
                instr_mask = torch.ones(instruction_max_len, dtype=torch.long)

            response_max_len = self.max_length // 2
            
            # Tokenize: 预留1个位置给EOS
            response_enc = self.tokenizer(
                response_text,
                truncation=True,
                max_length=response_max_len - 1, 
                add_special_tokens=False,
                return_tensors='pt'
            )
            resp_ids = response_enc['input_ids'].squeeze(0)
            
            eos_token = torch.tensor([self.tokenizer.eos_token_id], dtype=torch.long)
            resp_ids = torch.cat([resp_ids, eos_token])
            
            # Response 手动 Padding
            curr_resp_len = resp_ids.size(0)
            if curr_resp_len < response_max_len:
                padding = torch.full((response_max_len - curr_resp_len,), pad_token_id, dtype=torch.long)
                resp_ids = torch.cat([resp_ids, padding])
                resp_mask = torch.cat([torch.ones(curr_resp_len, dtype=torch.long), torch.zeros(response_max_len - curr_resp_len, dtype=torch.long)])
            else:
                resp_mask = torch.ones(response_max_len, dtype=torch.long)

            result = {
                'instruction': instr_ids,
                'response': resp_ids,
                'instruction_mask': instr_mask,
                'response_mask': resp_mask,
                'task': sample.get('_source', 'unknown'),
                'modality_data': None
            }
            
            if config.get('type') == 'multimodal_instruction' and 'image' in sample:
                try:
                    image = sample['image']
                    if isinstance(image, Image.Image):
                        image = image.convert('RGB')
                        image_tensor = image_transform(image)
                        result['modality_data'] = {'image': image_tensor}
                except Exception as e:
                    logger.debug(f"Error processing image: {e}")
            
            return result
        
        except Exception as e:
            logger.debug(f"Error getting item at index {idx}: {e}")
            import traceback
            traceback.print_exc()
            return None


class PreferenceDataset(Dataset):
    def __init__(
        self,
        dataset_name: str = 'hh_rlhf',
        tokenizer=None,
        max_length: int = 1024,
        max_samples: Optional[int] = None,
        split: str = 'train'
    ):
        super().__init__()
        
        if tokenizer is None:
            raise ValueError("tokenizer cannot be None")
        
        self.tokenizer = tokenizer
        self.max_length = max_length
        
        if dataset_name not in POSTTRAIN_DATASETS:
            raise ValueError(f"Unknown dataset: {dataset_name}. Available: {list(POSTTRAIN_DATASETS.keys())}")
        
        config = POSTTRAIN_DATASETS[dataset_name]
        if config.get('type') != 'preference':
            raise ValueError(f"{dataset_name} is not a preference dataset (type: {config.get('type')})")
        
        logger.info(f"Loading preference dataset: {dataset_name}")
        
        load_kwargs = {
            'path': config['hf_path'],
            'split': split,
            'cache_dir': HF_CACHE_DIR,
        }
        
        if 'config' in config:
            load_kwargs['name'] = config['config']
        
        self.dataset = load_dataset(**load_kwargs)
        
        self.chosen_field = config.get('chosen_field', 'chosen')
        self.rejected_field = config.get('rejected_field', 'rejected')
        
        if max_samples and len(self.dataset) > max_samples:
            self.dataset = self.dataset.select(range(max_samples))
        
        logger.info(f"Loaded {len(self.dataset)} preference pairs")

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, idx):
        try:
            sample = self.dataset[idx]
            
            chosen_text = sample.get(self.chosen_field, '')
            rejected_text = sample.get(self.rejected_field, '')
            
            if not chosen_text or not rejected_text:
                return None
            
            # Tokenize
            chosen_enc = self.tokenizer(
                chosen_text,
                max_length=self.max_length,
                truncation=True,
                padding='max_length',
                return_tensors='pt'
            )
            
            rejected_enc = self.tokenizer(
                rejected_text,
                max_length=self.max_length,
                truncation=True,
                padding='max_length',
                return_tensors='pt'
            )
            
            return (
                chosen_enc['input_ids'].squeeze(0),
                rejected_enc['input_ids'].squeeze(0),
                chosen_enc['attention_mask'].squeeze(0),
                rejected_enc['attention_mask'].squeeze(0)
            )
        
        except Exception as e:
            logger.debug(f"Error getting preference item at index {idx}: {e}")
            return None


def collate_fn_v2(batch):
    batch = [item for item in batch if item is not None]
    
    if not batch:
        logger.warning("Empty batch after filtering None values")
        # 返回一个空的占位batch而不是None
        return {
            'input_ids': torch.empty(0),
            'attention_mask': torch.empty(0)
        }
    
    # 检查是否是preference数据
    if isinstance(batch[0], tuple):
        if len(batch[0]) == 4:  # 包含attention_mask
            chosen = torch.stack([item[0] for item in batch])
            rejected = torch.stack([item[1] for item in batch])
            chosen_mask = torch.stack([item[2] for item in batch])
            rejected_mask = torch.stack([item[3] for item in batch])
            return {
                'chosen': chosen,
                'rejected': rejected,
                'chosen_mask': chosen_mask,
                'rejected_mask': rejected_mask
            }
        else: 
            chosen = torch.stack([item[0] for item in batch])
            rejected = torch.stack([item[1] for item in batch])
            return {'chosen': chosen, 'rejected': rejected}

    keys = batch[0].keys()
    collated = {}

    for key in keys:
        if key in ['instruction', 'response', 'instruction_mask', 
                   'response_mask', 'input_ids', 'attention_mask']:
            tensors = [item[key] for item in batch if item.get(key) is not None]
            if tensors:
                collated[key] = torch.stack(tensors)
            else:
                collated[key] = None
        elif key == 'modality_data':
            # 处理多模态数据
            modality_list = [item[key] for item in batch if item.get(key) is not None]
            if modality_list and any(m is not None for m in modality_list):
                # 收集图像
                images = [m.get('image') for m in modality_list if m and 'image' in m]
                if images:
                    collated[key] = {'image': torch.stack(images)}
                else:
                    collated[key] = None
            else:
                collated[key] = None
        else:
            collated[key] = [item[key] for item in batch]

    return collated


def create_pretrain_dataloader(
    mix_name: str = 'default',
    tokenizer=None,
    batch_size: int = 8,
    num_workers: int = 4,
    max_length: int = 2048,
    max_samples: Optional[int] = None
):
    dataset = PreTrainDataset(
        mix_name=mix_name,
        tokenizer=tokenizer,
        max_length=max_length,
        streaming=True,
        max_samples=max_samples
    )
    return DataLoader(
        dataset,
        batch_size=batch_size,
        num_workers=num_workers,
        collate_fn=collate_fn_v2
    )


def create_posttrain_dataloader(
    mix_name: str = 'default',
    tokenizer=None,
    batch_size: int = 8,
    num_workers: int = 4,
    max_length: int = 2048,
    max_samples: Optional[int] = None,
    split: str = 'train',
    shuffle: bool = True
):
    dataset = PostTrainDataset(
        mix_name=mix_name,
        tokenizer=tokenizer,
        max_length=max_length,
        max_samples=max_samples,
        split=split
    )
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        num_workers=num_workers,
        collate_fn=collate_fn_v2,
        pin_memory=True,
        drop_last=False 
    )


def create_preference_dataloader(
    dataset_name: str = 'hh_rlhf',
    tokenizer=None,
    batch_size: int = 8,
    num_workers: int = 4,
    max_length: int = 1024,
    max_samples: Optional[int] = None,
    split: str = 'train',
    shuffle: bool = True
):
    dataset = PreferenceDataset(
        dataset_name=dataset_name,
        tokenizer=tokenizer,
        max_length=max_length,
        max_samples=max_samples,
        split=split
    )
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        num_workers=num_workers,
        collate_fn=collate_fn_v2,
        pin_memory=True
    )