Update grpo_dataloader.py

grpo_dataloader.py

@@ -1,6 +1,3 @@
-"""
-GRPO专用数据加载器
-"""
 import torch
 from torch.utils.data import Dataset, DataLoader
 from datasets import load_dataset, interleave_datasets
@@ -18,9 +15,6 @@ from data_config import (
 
 
 class GRPOPromptDataset(Dataset):
-    """
-    GRPO Prompt数据集 - 用于生成阶段
-    """
     def __init__(
         self,
         mix_name: str = 'default',
@@ -35,8 +29,7 @@ class GRPOPromptDataset(Dataset):
         
         self.tokenizer = tokenizer
         self.max_length = max_length
-        
-        # 获取混合配置
+
         if mix_name not in GRPO_PROMPT_MIX:
             raise ValueError(
                 f"Unknown mix: {mix_name}. "
@@ -46,12 +39,7 @@ class GRPOPromptDataset(Dataset):
         mix_config = GRPO_PROMPT_MIX[mix_name]
         dataset_names = mix_config.get('datasets', [])
         weights = mix_config.get('weights', [])
-        
-        logger.info(f"Loading GRPO prompt mix: {mix_name}")
-        logger.info(f"  Datasets: {dataset_names}")
-        logger.info(f"  Weights: {weights}")
-        
-        # 加载数据集
+
         all_datasets = []
         
         for name in dataset_names:
@@ -60,12 +48,10 @@ class GRPOPromptDataset(Dataset):
                 continue
             
             config = GRPO_DATASETS[name]
-            
-            # 验证文件存在
+
             data_file = config.get('data_files')
             if data_file and not os.path.exists(data_file):
                 logger.error(f"Data file not found: {data_file}")
-                logger.error(f"请先运行 download_grpo_datasets.py 下载数据")
                 continue
             
             try:
@@ -79,13 +65,10 @@ class GRPOPromptDataset(Dataset):
                     load_kwargs['data_files'] = config['data_files']
                 
                 ds = load_dataset(**load_kwargs)
-                
-                # 限制样本数
                 if config.get('max_samples'):
                     ds = ds.select(range(min(len(ds), config['max_samples'])))
                 
                 all_datasets.append(ds)
-                logger.info(f"  Loaded {name}: {len(ds)} samples")
                 
             except Exception as e:
                 logger.error(f"Error loading {name}: {e}")
@@ -93,8 +76,7 @@ class GRPOPromptDataset(Dataset):
         
         if not all_datasets:
             raise ValueError("No datasets loaded successfully")
-        
-        # 合并数据集
+
         if len(all_datasets) == 1:
             self.dataset = all_datasets[0]
         else:
@@ -107,7 +89,6 @@ class GRPOPromptDataset(Dataset):
                 stopping_strategy='all_exhausted'
             )
         
-        # 限制总样本数
         if max_samples and len(self.dataset) > max_samples:
             self.dataset = self.dataset.select(range(max_samples))
         
@@ -119,15 +100,12 @@ class GRPOPromptDataset(Dataset):
     def __getitem__(self, idx):
         try:
             sample = self.dataset[idx]
-            
-            # 提取prompt
+
             prompt = sample.get('prompt', '')
             
             if not prompt:
                 logger.warning(f"Empty prompt at index {idx}")
                 return None
-            
-            # Tokenize (不添加EOS，因为这是prompt)
             encoding = self.tokenizer(
                 prompt,
                 max_length=self.max_length,
@@ -149,8 +127,6 @@ class GRPOPromptDataset(Dataset):
 
 
 def grpo_collate_fn(batch):
-    """GRPO专用collate函数"""
-    # 过滤None
     batch = [item for item in batch if item is not None]
     
     if not batch:
@@ -172,7 +148,6 @@ def create_grpo_prompt_dataloader(
     max_samples: Optional[int] = None,
     shuffle: bool = True
 ):
-    """创建GRPO prompt数据加载器"""
     dataset = GRPOPromptDataset(
         mix_name=mix_name,
         tokenizer=tokenizer,