apps_stats.py

import argparse
import matplotlib.pyplot as plt
import numpy as np
from datasets import load_dataset
import seaborn as sns
import json
import pandas as pd

def get_dataset_stats():
    """Generate and display statistics for the APPS dataset.

    This function loads the APPS dataset, computes various statistics, and
    visualizes them using matplotlib and seaborn. The statistics include:
    - Total number of samples per difficulty and split
    - Distribution of number of tests by difficulty and split
    """
    # Load APPS dataset
    dataset = load_dataset("codeparrot/apps", trust_remote_code=True)
    
    # Initialize data collection
    stats = {
        'difficulty': [],
        'num_tests': [],
        'has_solution': [],
        'num_solutions': [],
        'split': []
    }

    # Define consistent difficulty order
    difficulty_order = ['introductory', 'interview', 'competition']
    # Define consistent colors for each difficulty
    difficulty_colors = {
        'introductory': '#1f77b4',  # blue
        'interview': '#ff7f0e',     # orange  
        'competition': '#2ca02c'    # green
    }

    # Collect statistics
    for split in ['train', 'test']:
        for sample in dataset[split]:
            try:
                input_output = json.loads(sample['input_output']) if sample['input_output'].strip() else None
                num_tests = len(input_output["inputs"]) if input_output else 0
            except:
                num_tests = 0
                
            try:
                solutions = json.loads(sample['solutions'])
                num_solutions = len(solutions)
            except:
                solutions = []
                num_solutions = 0
                
            stats['difficulty'].append(sample['difficulty'])
            stats['num_tests'].append(num_tests)
            stats['has_solution'].append(num_solutions > 0)
            stats['num_solutions'].append(num_solutions)
            stats['split'].append(split)

    df = pd.DataFrame(stats)
    
    # Create figure with subplots
    fig = plt.figure(figsize=(20, 12))

    # 1. Bar plot of total samples per difficulty and split
    ax = plt.subplot(2, 2, 1)
    difficulty_counts = df.groupby(['split', 'difficulty']).size().unstack()[difficulty_order]
    difficulty_counts.plot(kind='bar', ax=ax, color=[difficulty_colors[d] for d in difficulty_order])
    ax.set_title('Total Number of Samples per Difficulty and Split')
    ax.set_xlabel('Split')
    ax.set_ylabel('Count')
    ax.legend(title='Difficulty')
    ax.legend(loc='upper right')

    # Print number of samples with 0 tests
    print("\nSamples with 0 tests:")
    print("-" * 50)
    zero_tests = len(df[df['num_tests'] == 0])
    print(f"Total samples with 0 tests: {zero_tests}")
    print("\nBreakdown by difficulty and split:")
    print(df[df['num_tests'] == 0].groupby(['split', 'difficulty']).size())
    
    
    # 2. Violin plot of number of tests by difficulty and split
    plt.subplot(2, 2, 2)
    # Filter to bottom 95% of data
    p95 = np.percentile(df['num_tests'], 95)
    df_filtered = df[df['num_tests'] <= p95]
    # Create violin plot with test on left, train on right
    sns.violinplot(data=df_filtered, x='difficulty', y='num_tests', hue='split', 
                  hue_order=['test', 'train'], split=True, 
                  order=difficulty_order)
    plt.title('Number of Tests Distribution by Difficulty and Split (Bottom 95%)')
    plt.xlabel('Difficulty')
    plt.ylabel('Number of Tests')
    plt.ylim(bottom=0)  # Set lower limit to 0
    
    
    # 3. Samples with/without solutions per difficulty and split
    ax = plt.subplot(2, 2, 3)
    solution_by_diff = pd.crosstab(df['split'], [df['difficulty'], df['has_solution']])
    
    # Set up bar positions
    bar_width = 0.25
    r1 = np.arange(len(solution_by_diff))
    r2 = [x + bar_width for x in r1]
    r3 = [x + bar_width for x in r2]
    
    # Plot bars for each difficulty with fill based on split
    ax.bar(r1, solution_by_diff[('introductory', False)], width=bar_width, label='Introductory - No Solution', 
           color=difficulty_colors['introductory'], fill=True)
    ax.bar(r1, solution_by_diff[('introductory', True)], width=bar_width, 
           bottom=solution_by_diff[('introductory', False)], label='Introductory - Has Solution',
           color=difficulty_colors['introductory'], fill=True, hatch='///')
    
    ax.bar(r2, solution_by_diff[('interview', False)], width=bar_width, label='Interview - No Solution',
           color=difficulty_colors['interview'], fill=True)
    ax.bar(r2, solution_by_diff[('interview', True)], width=bar_width,
           bottom=solution_by_diff[('interview', False)], label='Interview - Has Solution',
           color=difficulty_colors['interview'], fill=True, hatch='///')
    
    ax.bar(r3, solution_by_diff[('competition', False)], width=bar_width, label='Competition - No Solution',
           color=difficulty_colors['competition'], fill=True)
    ax.bar(r3, solution_by_diff[('competition', True)], width=bar_width,
           bottom=solution_by_diff[('competition', False)], label='Competition - Has Solution', 
           color=difficulty_colors['competition'], fill=True, hatch='///')
    
    ax.set_title('Samples With/Without Solutions by Split and Difficulty')
    ax.set_xlabel('Split')
    ax.set_ylabel('Count')
    ax.set_xticks([r + bar_width for r in range(len(solution_by_diff))])
    ax.set_xticklabels(solution_by_diff.index)
    ax.legend(loc='upper right')

    # 4. Samples with/without tests per difficulty and split
    plt.subplot(2, 2, 4)
    tests_by_diff = pd.crosstab(df['split'], [df['difficulty'], df['num_tests'] > 0])
    
    # Set up bar positions
    bar_width = 0.25
    r1 = np.arange(len(tests_by_diff))
    r2 = [x + bar_width for x in r1]
    r3 = [x + bar_width for x in r2]
    
    # Plot bars for each difficulty with fill based on split
    plt.bar(r1, tests_by_diff[('introductory', False)], width=bar_width, label='Introductory - No Tests', 
           color=difficulty_colors['introductory'], fill=True)
    plt.bar(r1, tests_by_diff[('introductory', True)], width=bar_width, 
           bottom=tests_by_diff[('introductory', False)], label='Introductory - Has Tests',
           color=difficulty_colors['introductory'], fill=True, hatch='///')
    
    plt.bar(r2, tests_by_diff[('interview', False)], width=bar_width, label='Interview - No Tests',
           color=difficulty_colors['interview'], fill=True)
    plt.bar(r2, tests_by_diff[('interview', True)], width=bar_width,
           bottom=tests_by_diff[('interview', False)], label='Interview - Has Tests',
           color=difficulty_colors['interview'], fill=True, hatch='///')
    
    plt.bar(r3, tests_by_diff[('competition', False)], width=bar_width, label='Competition - No Tests',
           color=difficulty_colors['competition'], fill=True)
    plt.bar(r3, tests_by_diff[('competition', True)], width=bar_width,
           bottom=tests_by_diff[('competition', False)], label='Competition - Has Tests', 
           color=difficulty_colors['competition'], fill=True, hatch='///')
    
    plt.title('Distribution of Samples With/Without Tests by Difficulty and Split')
    plt.xlabel('Split')
    plt.ylabel('Count')
    plt.legend(loc='center left', bbox_to_anchor=(0.05, 0.5))
    plt.xticks([r + bar_width for r in range(len(tests_by_diff))], tests_by_diff.index)

    
    # Print summary statistics
    print("\nSummary Statistics:")
    print("-" * 50)
    print(f"Total number of samples: {len(df)}")
    print("\nAverage number of tests by difficulty and split:")
    print(df.groupby(['split', 'difficulty'])['num_tests'].mean())
    print("\nMedian number of tests by difficulty and split:")
    print(df.groupby(['split', 'difficulty'])['num_tests'].median())
    print("\nSamples with solutions by difficulty and split:")
    print(df.groupby(['split', 'difficulty'])['has_solution'].value_counts())

    print("\nMedian number of solutions by difficulty and split:")
    print(df.groupby(['split', 'difficulty'])['num_solutions'].median())
    
    # Calculate counts and percentages of samples with >=5 soluоtions
    solutions_summary = df.groupby(['split', 'difficulty']).agg({
        'num_solutions': [
            ('total_count', 'size'),
            ('many_solutions_count', lambda x: sum(x >= 5)),
        ]
    })
    solutions_summary[('num_solutions', 'many_solutions_pct')] = (
        solutions_summary[('num_solutions', 'many_solutions_count')] / 
        solutions_summary[('num_solutions', 'total_count')] * 100
    ).round(1)
    
    print("\nSamples with >=5 solutions by difficulty and split:")
    print("Count and percentage:")
    print(solutions_summary['num_solutions'].to_string())

    # Calculate counts and percentages for >=3 solutions
    solutions_summary_3 = df.groupby(['split', 'difficulty']).agg({
        'num_solutions': [
            ('total_count', 'size'),
            ('many_solutions_count', lambda x: sum(x >= 3)),
        ]
    })
    solutions_summary_3[('num_solutions', 'many_solutions_pct')] = (
        solutions_summary_3[('num_solutions', 'many_solutions_count')] / 
        solutions_summary_3[('num_solutions', 'total_count')] * 100
    ).round(1)
    
    print("\nSamples with >=3 solutions by difficulty and split:")
    print("Count and percentage:")
    print(solutions_summary_3['num_solutions'].to_string())

    # Calculate counts and percentages for >=2 solutions
    solutions_summary_2 = df.groupby(['split', 'difficulty']).agg({
        'num_solutions': [
            ('total_count', 'size'),
            ('many_solutions_count', lambda x: sum(x >= 2)),
        ]
    })
    solutions_summary_2[('num_solutions', 'many_solutions_pct')] = (
        solutions_summary_2[('num_solutions', 'many_solutions_count')] / 
        solutions_summary_2[('num_solutions', 'total_count')] * 100
    ).round(1)
    
    print("\nSamples with >=2 solutions by difficulty and split:")
    print("Count and percentage:")
    print(solutions_summary_2['num_solutions'].to_string())
    
    # Adjust layout to prevent overlap
    plt.tight_layout()
    plt.savefig('apps_dataset_stats.png')
    plt.close()

if __name__ == "__main__":
    get_dataset_stats()