test_script.py

"""Quick start script for paper replication experiments.

This script helps you verify your setup and run a small test before
running the full experiments.
"""

import asyncio
import json
import logging
from pathlib import Path

from experiments.experiemnts import Experiments
from utils.data_loader import load_data, save_to_json

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


def check_data_files():
    """Check if required data files exist."""
    logger.info("Checking data files...")
    
    required_files = [
        "articles/cnn_train_articles.json",
        "articles/xsum_train_articles.json",
        "summaries/cnn_summaries_gpt_4.1_nano.json",
        "summaries/cnn_summaries_claude-3-haiku-20240307.json",
        "summaries/cnn_summaries_human.json",
        "summaries/xsum_summaries_gpt_4.1_nano.json",
        "summaries/xsum_summaries_claude-3-haiku-20240307.json",
        "summaries/xsum_summaries_human.json",
    ]
    
    missing = []
    for file_path in required_files:
        if not Path(file_path).exists():
            missing.append(file_path)
    
    if missing:
        logger.warning("Missing files:")
        for f in missing:
            logger.warning(f"  - {f}")
        return False
    
    logger.info("✓ All data files found!")
    return True


def check_data_format():
    """Check if data files are in correct format."""
    logger.info("Checking data format...")
    
    try:
        responses, articles, keys = load_data("cnn")
        
        logger.info(f"✓ Loaded {len(keys)} articles from CNN dataset")
        logger.info(f"✓ Found summaries from: {list(responses.keys())}")
        
        # Check if we have at least one complete sample
        if keys:
            key = keys[0]
            logger.info(f"\nSample article key: {key}")
            logger.info(f"Article preview: {articles[key][:100]}...")
            
            for source in responses:
                if key in responses[source]:
                    logger.info(f"{source} summary preview: {responses[source][key][:80]}...")
        
        return True
        
    except Exception as e:
        logger.error(f"Error loading data: {e}")
        return False


async def run_mini_experiment():
    """Run a mini experiment on a few samples to test the setup."""
    logger.info("\n" + "="*80)
    logger.info("Running mini experiment (first 3 samples)...")
    logger.info("="*80 + "\n")
    
    try:
        # Test with gpt-4.1-nano on CNN dataset
        exp = Experiments(
            model_name="gpt-4.1-nano",
            dataset="cnn",
            max_concurrent=2,
        )
        
        # Load data and limit to first 3 keys
        responses, articles, keys = load_data("cnn")
        test_keys = keys[:3]
        
        logger.info(f"Testing with {len(test_keys)} samples")
        logger.info(f"Keys: {test_keys}\n")
        
        # Monkey-patch the load_data function to return limited data
        original_load_data = exp.__class__.__module__
        
        # Run detection experiment on limited data
        results = []
        
        for i, key in enumerate(test_keys):
            logger.info(f"Processing sample {i+1}/{len(test_keys)}: {key}")
            
            article = articles[key]
            source_summary = responses["gpt-4.1-nano"][key]
            
            # Test with human summary
            if "human" in responses and key in responses["human"]:
                human_summary = responses["human"][key]
                
                logger.info("  Testing detection with human summary...")
                forward = await exp._get_model_choice_with_logprobs(
                    source_summary, human_summary, article, "detection"
                )
                
                if forward:
                    choice = forward[0]["token"]
                    prob = forward[0]["logprob"]
                    logger.info(f"    Forward choice: {choice} (logprob: {prob:.4f})")
                
                logger.info("  Testing comparison with human summary...")
                comparison = await exp._get_model_choice_with_logprobs(
                    source_summary, human_summary, article, "comparison"
                )
                
                if comparison:
                    choice = comparison[0]["token"]
                    prob = comparison[0]["logprob"]
                    logger.info(f"    Comparison choice: {choice} (logprob: {prob:.4f})")
        
        logger.info("\n✓ Mini experiment completed successfully!")
        logger.info("You can now run the full experiments with: python run_experiments.py")
        return True
        
    except Exception as e:
        logger.error(f"\n✗ Mini experiment failed: {e}", exc_info=True)
        return False


async def main():
    """Main function."""
    logger.info("="*80)
    logger.info("PAPER REPLICATION - QUICK START")
    logger.info("="*80 + "\n")
    
    # Step 1: Check data files
    if not check_data_files():
        logger.error("\n✗ Setup incomplete: Missing data files")
        logger.info("\nPlease ensure you have:")
        logger.info("1. Article files in articles/ directory")
        logger.info("2. Summary files in summaries/ directory")
        logger.info("\nSee PAPER_REPLICATION.md for details")
        return
    
    # Step 2: Check data format
    if not check_data_format():
        logger.error("\n✗ Setup incomplete: Data format issues")
        logger.info("\nPlease check:")
        logger.info("1. JSON files are valid")
        logger.info("2. Keys match across articles and summaries")
        logger.info("\nSee PAPER_REPLICATION.md for details")
        return
    
    # Step 3: Run mini experiment
    logger.info("\n" + "="*80)
    logger.info("All checks passed! Running mini experiment...")
    logger.info("="*80 + "\n")
    
    success = await run_mini_experiment()
    
    if success:
        logger.info("\n" + "="*80)
        logger.info("READY TO GO!")
        logger.info("="*80)
        logger.info("\nNext steps:")
        logger.info("1. Run full experiments: python run_paper_experiments.py")
        logger.info("2. Or run specific experiment: python run_paper_experiments.py gpt-4.1-nano cnn detection_comparison")
        logger.info("3. Analyze results: python analyze_paper_results.py")
        logger.info("\nSee PAPER_REPLICATION.md for more details")
    else:
        logger.error("\n" + "="*80)
        logger.error("SETUP ISSUES DETECTED")
        logger.error("="*80)
        logger.info("\nPlease check:")
        logger.info("1. API keys are configured (OPENAI_API_KEY, ANTHROPIC_API_KEY)")
        logger.info("2. Model names in config match data file names")
        logger.info("3. Internet connection is working")
        logger.info("\nSee PAPER_REPLICATION.md for troubleshooting")


if __name__ == "__main__":
    asyncio.run(main())