sample_data.py

"""
数据采样脚本 - 从完整的GSM混合数据中选择40万条训练数据
"""
import json
import random
import os

def sample_training_data(input_path: str, output_path: str, sample_size: int = 400000):
    print(f"Loading data from {input_path}...")
    with open(input_path, 'r', encoding='utf-8') as f:
        full_data = json.load(f)
    print(f"Original data size: {len(full_data)} examples")
    
    if len(full_data) <= sample_size:
        print(f"Warning: Requested {sample_size} samples, but only {len(full_data)} available.")
        print("Using all available data.")
        sampled_data = full_data
    else:
        random.seed(42)
        sampled_data = random.sample(full_data, sample_size)
    
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    
    print(f"Saving sampled data to {output_path}...")
    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(sampled_data, f, indent=2, ensure_ascii=False)
    
    print(f"Successfully created sampled dataset with {len(sampled_data)} examples")
    return output_path

def main():
    input_train_path = "./data/hybrid_dataset_train.json"
    input_val_path = "./data/hybrid_dataset_val.json"
    
    output_train_path = "./data/hybrid_dataset_train_400k.json"
    output_val_path = "./data/hybrid_dataset_val_sampled.json"
    if not os.path.exists(input_train_path):
        print(f"Error: Training data file not found: {input_train_path}")
        print("Please run the data generator first.")
        return
    
    if not os.path.exists(input_val_path):
        print(f"Error: Validation data file not found: {input_val_path}")
        print("Please run the data generator first.")
        return
 
    sample_training_data(input_train_path, output_train_path, sample_size=400000)
    sample_training_data(input_val_path, output_val_path, sample_size=5000)

    print("\n" + "="*60)
    print("Data Sampling Summary")
    print("="*60)
    with open(output_train_path, 'r') as f:
        train_data = json.load(f)
    with open(output_val_path, 'r') as f:
        val_data = json.load(f)
    
    print(f"Training data: {len(train_data):,} examples")
    print(f"Validation data: {len(val_data):,} examples")
    print(f"Total: {len(train_data) + len(val_data):,} examples")
 
    # print(f"\nSample training examples:")
    # for i, example in enumerate(train_data[:2]):
    #     print(f"\nExample {i+1}:")
    #     print("-" * 40)
    #     # 只显示问题部分，不显示完整文本（太长）
    #     question = example.get('question', 'No question found')
    #     answer = example.get('answer', 'No answer found')
    #     print(f"Question: {question[:100]}...")
    #     print(f"Answer: {answer}")
    #     print(f"Type: {example.get('type', 'Unknown')}")

if __name__ == "__main__":
    main()