path: root/enhance_gee_processor.py

#!/usr/bin/env python3
"""
增强GEE处理器以支持真实数据集
支持Numina数学推理数据和其他真实数据源
"""
import pandas as pd
import numpy as np
from pathlib import Path
import json
import re
import sys
sys.path.append('.')

from dataset.gee_processor import GEEProcessor

class EnhancedGEEProcessor(GEEProcessor):
    """增强版GEE处理器，支持多种真实数据源"""
    
    def __init__(self, tokenizer):
        super().__init__(tokenizer)
        self.name_patterns = {
            'male': ['Tom', 'John', 'Mike', 'Bob', 'David', 'James', 'Robert', 'Michael', 'William', 'Richard'],
            'female': ['Sarah', 'Lisa', 'Emma', 'Alice', 'Mary', 'Jennifer', 'Linda', 'Elizabeth', 'Barbara', 'Susan']
        }
    
    def process_numina_data(self, file_path: str, target_size: int = 1000) -> list:
        """处理Numina数学推理数据"""
        print(f"📊 处理Numina数据: {file_path}")
        
        # 读取parquet文件
        df = pd.read_parquet(file_path)
        print(f"原始数据量: {len(df)}")
        
        # 随机采样
        if len(df) > target_size:
            df = df.sample(n=target_size, random_state=42)
            print(f"采样后数据量: {len(df)}")
        
        processed_data = []
        for idx, row in df.iterrows():
            # 提取问题和答案
            problem = row.get('problem', row.get('question', ''))
            solution = row.get('solution', row.get('answer', ''))
            
            if problem and solution:
                # 生成性别平衡的变体
                male_version = self._genderize_text(problem, 'male')
                female_version = self._genderize_text(problem, 'female')
                
                processed_data.extend([
                    {
                        'input': self.apply_chat_template(male_version),
                        'output': solution,
                        'gender': 'male',
                        'original_id': idx,
                        'source': 'numina'
                    },
                    {
                        'input': self.apply_chat_template(female_version), 
                        'output': solution,
                        'gender': 'female',
                        'original_id': idx,
                        'source': 'numina'
                    }
                ])
        
        print(f"✅ 处理完成，生成 {len(processed_data)} 个样本")
        return processed_data
    
    def process_1shot_rlvr_data(self, file_path: str) -> list:
        """处理1shot RLVR数据"""
        print(f"�� 处理1shot RLVR数据: {file_path}")
        
        df = pd.read_parquet(file_path)
        print(f"原始数据量: {len(df)}")
        
        processed_data = []
        for idx, row in df.iterrows():
            # 根据实际数据结构调整
            prompt = row.get('prompt', row.get('input', ''))
            
            if prompt:
                # 生成性别变体
                for gender in ['male', 'female']:
                    genderized_prompt = self._genderize_text(prompt, gender)
                    
                    processed_data.append({
                        'input': self.apply_chat_template(genderized_prompt),
                        'gender': gender,
                        'original_id': idx,
                        'source': '1shot_rlvr'
                    })
        
        print(f"✅ 处理完成，生成 {len(processed_data)} 个样本")
        return processed_data
    
    def _genderize_text(self, text: str, target_gender: str) -> str:
        """将文本中的性别引用转换为指定性别"""
        
        # 选择名字
        names = self.name_patterns[target_gender]
        
        # 替换通用占位符
        if '[NAME]' in text or '{name}' in text:
            name = np.random.choice(names)
            text = text.replace('[NAME]', name).replace('{name}', name)
            return text
        
        # 检测现有性别名字并替换
        all_male_names = self.name_patterns['male']
        all_female_names = self.name_patterns['female'] 
        
        for male_name in all_male_names:
            if male_name in text:
                replacement = np.random.choice(names)
                text = text.replace(male_name, replacement)
                break
                
        for female_name in all_female_names:
            if female_name in text:
                replacement = np.random.choice(names)
                text = text.replace(female_name, replacement)
                break
        
        # 如果没有找到名字，随机添加一个
        if not any(name in text for name in all_male_names + all_female_names):
            name = np.random.choice(names)
            # 在合适的地方插入名字
            if "person" in text.lower():
                text = text.replace("person", name)
            elif "student" in text.lower():
                text = text.replace("student", f"student named {name}")
            elif "someone" in text.lower():
                text = text.replace("someone", name)
            else:
                # 在句子开头添加
                text = f"{name} is working on this problem: {text}"
        
        return text
    
    def create_balanced_dataset(self, data_sources: list, balance_method: str = 'oversample') -> list:
        """创建性别平衡的数据集"""
        
        all_data = []
        for source_config in data_sources:
            source_type = source_config['type']
            file_path = source_config['path']
            
            if source_type == 'numina':
                data = self.process_numina_data(file_path, source_config.get('target_size', 1000))
            elif source_type == '1shot_rlvr':
                data = self.process_1shot_rlvr_data(file_path)
            else:
                print(f"⚠️ 未知数据源类型: {source_type}")
                continue
                
            all_data.extend(data)
        
        # 统计性别分布
        male_data = [item for item in all_data if item['gender'] == 'male']
        female_data = [item for item in all_data if item['gender'] == 'female']
        
        print(f"\n📊 原始数据分布:")
        print(f"   男性样本: {len(male_data)}")
        print(f"   女性样本: {len(female_data)}")
        
        # 平衡处理
        if balance_method == 'oversample':
            target_size = max(len(male_data), len(female_data))
            
            if len(male_data) < target_size:
                male_data = male_data * (target_size // len(male_data)) + male_data[:target_size % len(male_data)]
            if len(female_data) < target_size:
                female_data = female_data * (target_size // len(female_data)) + female_data[:target_size % len(female_data)]
                
        elif balance_method == 'undersample':
            target_size = min(len(male_data), len(female_data))
            male_data = male_data[:target_size]
            female_data = female_data[:target_size]
        
        balanced_data = male_data + female_data
        np.random.shuffle(balanced_data)
        
        print(f"📊 平衡后数据分布:")
        male_count = sum(1 for item in balanced_data if item['gender'] == 'male')
        female_count = sum(1 for item in balanced_data if item['gender'] == 'female')
        print(f"   男性样本: {male_count}")
        print(f"   女性样本: {female_count}")
        print(f"   总样本数: {len(balanced_data)}")
        
        return balanced_data

def main():
    """示例用法"""
    from transformers import AutoTokenizer
    
    print("🔧 测试增强版GEE处理器...")
    
    # 初始化
    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Math-1.5B-Instruct", trust_remote_code=True)
    processor = EnhancedGEEProcessor(tokenizer)
    
    # 配置数据源
    data_sources = [
        {
            'type': 'numina',
            'path': 'dataset/numina/numina_00.parquet',
            'target_size': 100  # 测试用小样本
        }
        # 可以添加更多数据源
    ]
    
    # 处理数据
    try:
        balanced_data = processor.create_balanced_dataset(data_sources, balance_method='oversample')
        
        # 保存结果
        output_file = 'enhanced_training_data.json'
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(balanced_data, f, indent=2, ensure_ascii=False)
        
        print(f"✅ 增强数据已保存: {output_file}")
        
        # 显示示例
        print(f"\n📝 数据示例:")
        for i, sample in enumerate(balanced_data[:4]):
            print(f"  示例 {i+1} ({sample['gender']}):")
            print(f"    输入: {sample['input'][:100]}...")
            print()
            
    except Exception as e:
        print(f"❌ 处理失败: {e}")

if __name__ == "__main__":
    main()