src/utils/data_utils.py

import os
from datasets import load_dataset
from transformers import AutoTokenizer
import torch
from torch.utils.data import DataLoader

def custom_collate(batch):
    return {
        'src_ids': torch.stack([torch.tensor(x['src_ids']) for x in batch]),
        'src_mask': torch.stack([torch.tensor(x['src_mask']) for x in batch]),
        'tgt_ids': torch.stack([torch.tensor(x['tgt_ids']) for x in batch]),
        'tgt_mask': torch.stack([torch.tensor(x['tgt_mask']) for x in batch]),
        # 保留测试用例用于验证 (仅 Eval 时有效)
        'test_code': [x.get('test_code', "") for x in batch],
        'entry_point': [x.get('entry_point', "") for x in batch]
    }

def prepare_data(task_name, tokenizer, max_len, batch_size, split="train"):
    """
    支持 split 参数，方便划分训练集和测试集
    """
    print(f"Loading {task_name} ({split})...")
    
    if task_name == "codexglue":
        # 训练集：Microsoft CodeXGLUE (Python Refinement)
        # 包含 GitHub Bug -> Fix
        dataset = load_dataset("./code_x_glue_cc_code_refinement_full", "medium", split=split)

        # 40k
        if split == "train": dataset = dataset.select(range(40000)) 

        # Case A: 标准修复数据 (有 source 和 target)
        if 'source' in cols and 'target' in cols:
            print(">> Detected standard refinement pairs.")
            def preprocess_standard(ex):
                src = tokenizer(ex['source'], max_length=max_len, padding="max_length", truncation=True)
                tgt = tokenizer(ex['target'], max_length=max_len, padding="max_length", truncation=True)
                return {
                    'src_ids': src['input_ids'], 'src_mask': src['attention_mask'],
                    'tgt_ids': tgt['input_ids'], 'tgt_mask': tgt['attention_mask']
                }
            preprocess_fn = preprocess_standard

        # Case B: 只有代码 (有 code)，需要人工注入 Bug
        elif 'code' in cols:
            print(">> Detected raw code. Not to inject synthetic bugs...")
            
        else:
            raise ValueError(f"Dataset columns {cols} not recognized. Need 'source'/'target' or 'code'.")
        
        def preprocess(ex):
            buggy = ex['source']
            fixed = ex['target']
            
            src = tokenizer(buggy, max_length=max_len, padding="max_length", truncation=True)
            tgt = tokenizer(fixed, max_length=max_len, padding="max_length", truncation=True)
            
            return {
                'src_ids': src['input_ids'], 'src_mask': src['attention_mask'],
                'tgt_ids': tgt['input_ids'], 'tgt_mask': tgt['attention_mask']
            }
        
        # 移除原始列
        cols = dataset.column_names
    
    elif task_name == "humanevalpack":
        # 验证集：HumanEvalPack (Fix Task)
        # 包含 Buggy Code 和 对应的 Unit Tests
        dataset = load_dataset("./bigcode_humanevalpack_full", "python", split="test") # 只有 test 集
        
        # 筛选出 FIX 任务
        dataset = dataset.filter(lambda x: x['task_id'].startswith("Python/FIX"))
        
        def preprocess(ex):
            # prompt 是前面的描述，buggy_solution 是有 bug 的代码
            # 为了简化，我们把 prompt + buggy_solution 作为输入
            full_buggy = ex['prompt'] + "\n" + ex['buggy_solution']
            full_fixed = ex['prompt'] + "\n" + ex['canonical_solution']
            
            src = tokenizer(full_buggy, max_length=max_len, padding="max_length", truncation=True)
            tgt = tokenizer(full_fixed, max_length=max_len, padding="max_length", truncation=True)
            
            return {
                'src_ids': src['input_ids'], 'src_mask': src['attention_mask'],
                'tgt_ids': tgt['input_ids'], 'tgt_mask': tgt['attention_mask'],
                'test_code': ex['test'],         # 核心：保留测试代码
                'entry_point': ex['entry_point'] # 核心：保留入口函数名
            }
        
        # 保留所有列用于 debug，dataset.map 会自动处理返回的 dict
        cols = [] # 不自动删除列，我们需要 test 列在 collate 中处理
    
    # --- 1. Load Dataset ---
    elif task_name == "wiki":
         # 尝试本地加载，失败则下载
        try:
            dataset = load_dataset("./wikilarge-dataset")
        except:
            print("Local load failed, downloading from Hub...")
            dataset = load_dataset("wikilarge")
        
        # 手动划分: train用前10000条, test用后1000条 (做demo够了，全量太慢)
        if split == "train":
            dataset = dataset['train'].select(range(20000))
        else:
            # 假设总共有 ~290k，我们取后面一点做测试
            dataset = dataset['train'].select(range(20000, 25000))

        # 自动探测列名
        cols = dataset.column_names
        print(f"Wiki Dataset Columns: {cols}")
        
        # 映射列名到 src/tgt
        if 'src' in cols and 'dst' in cols:
            src_key, tgt_key = 'src', 'dst'
        elif 'Normal' in cols and 'Simple' in cols:
            src_key, tgt_key = 'Normal', 'Simple'
        else:
            raise ValueError(f"Unknown column format for WikiLarge: {cols}")

        def preprocess(ex):
            # Source (Complex) -> Target (Simple)
            src = tokenizer(ex[src_key], max_length=max_len, padding="max_length", truncation=True)
            tgt = tokenizer(ex[tgt_key], max_length=max_len, padding="max_length", truncation=True)
            return {
                'src_ids': src['input_ids'], 'src_mask': src['attention_mask'],
                'tgt_ids': tgt['input_ids'], 'tgt_mask': tgt['attention_mask']
            }
        
    elif task_name == "mbpp":
        dataset = load_dataset("mbpp", split="train[:500]")
        print(f"MBPP Dataset Columns: {dataset.column_names}")
        
        # MBPP 自重建任务: src=code, tgt=code
        def preprocess(ex):
            enc = tokenizer(ex['code'], max_length=max_len, padding="max_length", truncation=True)
            return {
                'src_ids': enc['input_ids'], 'src_mask': enc['attention_mask'],
                'tgt_ids': enc['input_ids'], 'tgt_mask': enc['attention_mask']
            }
    
    else:
        raise ValueError(f"Unknown task: {task_name}")

    # --- 2. Map & Batch ---
    print(f"Preprocessing {task_name} data...")
    # 使用 remove_columns=dataset.column_names 确保删除所有原始列
    print(f"Preprocessing {len(dataset)} examples...")
    dataset = dataset.map(
        preprocess, 
        batched=True, 
        remove_columns=dataset.column_names, 
        num_proc=4
    )
    
    # Test 集不 shuffle，方便对齐
    return DataLoader(dataset, batch_size=batch_size, shuffle=(split=="train"), collate_fn=custom_collate)