path: root/smart_balanced_dataloader.py

#!/usr/bin/env python3
"""
智能平衡数据加载器 - 自动检测并重新生成不平衡批次
"""
import torch
from torch.utils.data import Dataset, DataLoader
import random
from typing import List, Dict
import warnings

class SmartBalancedGEEDataset(Dataset):
    def __init__(self, data: List[Dict]):
        self.data = data
        # 按性别分组
        self.male_data = [item for item in data if item['gender'] == 'male']
        self.female_data = [item for item in data if item['gender'] == 'female']
        
        print(f"📊 数据分布: male={len(self.male_data)}, female={len(self.female_data)}")
        
        # 确保有足够的数据
        if len(self.male_data) == 0 or len(self.female_data) == 0:
            raise ValueError("数据中必须包含男性和女性样本")
        
        # 确保有足够的样本进行重新生成
        min_samples = min(len(self.male_data), len(self.female_data))
        if min_samples < 10:
            warnings.warn(f"样本数量较少 (min={min_samples})，可能影响批次生成质量")
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        return self.data[idx]
    
    def generate_balanced_batch(self, batch_size: int, max_retries: int = 10):
        """生成一个平衡的批次，如果失败会重试"""
        if batch_size < 2:
            raise ValueError("batch_size必须>=2才能保证性别平衡")
        
        # 计算每个性别需要的样本数
        male_per_batch = batch_size // 2
        female_per_batch = batch_size - male_per_batch
        
        for attempt in range(max_retries):
            try:
                batch = []
                
                # 🔧 修复: 使用更好的随机策略确保平衡
                # 强制选择所需数量的男性样本
                if len(self.male_data) >= male_per_batch:
                    male_samples = random.sample(self.male_data, male_per_batch)
                else:
                    # 如果男性样本不够，用替换的方式采样
                    male_samples = random.choices(self.male_data, k=male_per_batch)
                
                batch.extend(male_samples)
                
                # 强制选择所需数量的女性样本
                if len(self.female_data) >= female_per_batch:
                    female_samples = random.sample(self.female_data, female_per_batch)
                else:
                    # 如果女性样本不够，用替换的方式采样
                    female_samples = random.choices(self.female_data, k=female_per_batch)
                
                batch.extend(female_samples)
                
                # 打乱批次内的顺序
                random.shuffle(batch)
                
                # 最终验证批次平衡性
                male_count = sum(1 for item in batch if item['gender'] == 'male')
                female_count = sum(1 for item in batch if item['gender'] == 'female')
                
                if male_count > 0 and female_count > 0:
                    return batch
                else:
                    print(f"❌ 尝试 {attempt+1}: 批次不平衡 (male={male_count}, female={female_count})，重新生成...")
                    
            except Exception as e:
                print(f"⚠️ 尝试 {attempt+1} 失败: {e}")
                continue
        
        # 如果所有尝试都失败，返回一个强制平衡的批次
        print(f"❌ {max_retries} 次尝试后仍然失败，强制生成平衡批次")
        return self._force_balanced_batch(batch_size)
    
    def _force_balanced_batch(self, batch_size: int):
        """强制生成一个平衡批次"""
        batch = []
        male_per_batch = batch_size // 2
        female_per_batch = batch_size - male_per_batch
        
        # 强制添加男性样本（允许重复）
        for _ in range(male_per_batch):
            batch.append(random.choice(self.male_data))
        
        # 强制添加女性样本（允许重复）
        for _ in range(female_per_batch):
            batch.append(random.choice(self.female_data))
        
        random.shuffle(batch)
        return batch

class SmartBalancedDataLoader:
    """智能平衡数据加载器"""
    def __init__(self, dataset: SmartBalancedGEEDataset, batch_size: int, num_batches: int):
        self.dataset = dataset
        self.batch_size = batch_size
        self.num_batches = num_batches
        self.current_idx = 0
        
        print(f"🧠 智能数据加载器初始化: batch_size={batch_size}, num_batches={num_batches}")
    
    def __iter__(self):
        self.current_idx = 0
        return self
    
    def __next__(self):
        if self.current_idx >= self.num_batches:
            raise StopIteration
        
        # 🔧 特殊处理第一个批次，确保绝对平衡
        if self.current_idx == 0:
            print(f"🎯 特殊处理第一个批次")
            batch = self._generate_guaranteed_balanced_batch()
        else:
            # 其他批次使用标准方法
            batch = self.dataset.generate_balanced_batch(self.batch_size)
        
        self.current_idx += 1
        
        # 应用collate函数并验证
        return self._smart_collate(batch)
    
    def _generate_guaranteed_balanced_batch(self):
        """保证生成平衡的第一个批次"""
        batch_size = self.batch_size
        male_per_batch = batch_size // 2
        female_per_batch = batch_size - male_per_batch
        
        batch = []
        
        # 强制选择男性样本（取前N个或轮换选择）
        if len(self.dataset.male_data) >= male_per_batch:
            # 不使用随机，而是轮换选择
            male_samples = self.dataset.male_data[:male_per_batch]
        else:
            # 重复选择
            male_samples = (self.dataset.male_data * ((male_per_batch // len(self.dataset.male_data)) + 1))[:male_per_batch]
        
        batch.extend(male_samples)
        
        # 强制选择女性样本（取前N个或轮换选择）
        if len(self.dataset.female_data) >= female_per_batch:
            # 不使用随机，而是轮换选择
            female_samples = self.dataset.female_data[:female_per_batch]
        else:
            # 重复选择
            female_samples = (self.dataset.female_data * ((female_per_batch // len(self.dataset.female_data)) + 1))[:female_per_batch]
        
        batch.extend(female_samples)
        
        # 验证
        male_count = sum(1 for item in batch if item['gender'] == 'male')
        female_count = sum(1 for item in batch if item['gender'] == 'female')
        print(f"🎯 第一批次: male={male_count}, female={female_count} (强制平衡)")
        
        # 打乱顺序
        random.shuffle(batch)
        
        return batch
    
    def _smart_collate(self, batch, max_regenerate: int = 3):
        """智能collate函数，如果检测到不平衡会重新生成"""
        inputs = [item["input"] for item in batch]
        genders = [item["gender"] for item in batch]
        
        # 检查批次平衡性
        male_count = sum(1 for g in genders if g == 'male')
        female_count = sum(1 for g in genders if g == 'female')
        
        # 如果不平衡，尝试重新生成
        regenerate_count = 0
        while (male_count == 0 or female_count == 0) and regenerate_count < max_regenerate:
            print(f"🔄 检测到不平衡批次 (male={male_count}, female={female_count})，重新生成...")
            
            # 重新生成批次
            batch = self.dataset.generate_balanced_batch(self.batch_size)
            inputs = [item["input"] for item in batch]
            genders = [item["gender"] for item in batch]
            
            # 重新检查
            male_count = sum(1 for g in genders if g == 'male')
            female_count = sum(1 for g in genders if g == 'female')
            regenerate_count += 1
        
        # 最终检查
        if male_count == 0:
            print("❌ 最终警告: 批次中没有男性样本！")
        if female_count == 0:
            print("❌ 最终警告: 批次中没有女性样本！")
        
        if male_count > 0 and female_count > 0:
            print(f"✅ 平衡批次: male={male_count}, female={female_count}")
        
        return {
            "input": inputs,
            "gender": genders
        }

def create_smart_balanced_dataloader(data: List[Dict], batch_size: int, num_batches: int = 10):
    """创建智能平衡数据加载器"""
    
    if batch_size < 2:
        print("⚠️ 警告: batch_size < 2，无法保证性别平衡")
        # 回退到普通DataLoader
        dataset = SmartBalancedGEEDataset(data)
        return DataLoader(dataset, batch_size=batch_size, shuffle=True)
    
    dataset = SmartBalancedGEEDataset(data)
    
    print(f"🧠 创建智能平衡数据加载器")
    print(f"   批次大小: {batch_size}")
    print(f"   批次数量: {num_batches}")
    print(f"   每批次配置: male={batch_size//2}, female={batch_size-batch_size//2}")
    
    return SmartBalancedDataLoader(dataset, batch_size, num_batches)

# 测试函数
if __name__ == "__main__":
    # 测试智能平衡数据加载器
    import sys
    sys.path.append('.')
    from dataset.gee_processor import GEEProcessor
    
    class MockTokenizer:
        def apply_chat_template(self, messages, tokenize=False, add_generation_prompt=True):
            return messages[0]["content"]
    
    processor = GEEProcessor(MockTokenizer())
    test_data = processor.create_test_data(num_samples=20)
    
    print("🧪 测试智能平衡数据加载器")
    dataloader = create_smart_balanced_dataloader(test_data, batch_size=4, num_batches=5)
    
    for i, batch in enumerate(dataloader):
        print(f"\n=== 批次 {i+1} ===")
        print(f"输入数量: {len(batch['input'])}")
        print(f"性别分布: {batch['gender']}")
        
        # 验证平衡性
        male_count = sum(1 for g in batch['gender'] if g == 'male')
        female_count = sum(1 for g in batch['gender'] if g == 'female')
        
        if male_count > 0 and female_count > 0:
            print(f"✅ 批次完美平衡: male={male_count}, female={female_count}")
        else:
            print(f"❌ 批次仍然不平衡: male={male_count}, female={female_count}")
        
        if i >= 4:  # 测试5个批次
            break
    
    print("\n�� 智能平衡数据加载器测试完成!")