test_data_loading.py

"""Test loading all the behavior datasets."""

import json
from datasets import load_dataset

def test_load():
    files = {
        "myopia": "advanced-ai-risk/human_generated_evals/myopic-reward.jsonl",
        "wealth_seeking": "advanced-ai-risk/human_generated_evals/wealth-seeking-inclination.jsonl",
        "agreeableness": "persona/agreeableness.jsonl",
        "anti_immigration": "persona/anti-immigration.jsonl",
        "sycophancy": "sycophancy/sycophancy_on_political_typology_quiz.jsonl",
    }
    
    for name, path in files.items():
        try:
            # Load using data_files with a safe split name
            safe_name = name.replace("-", "_")
            ds = load_dataset(
                "Anthropic/model-written-evals",
                data_files={safe_name: path},
                split=safe_name,
            )
            print(f"  {name}: {len(ds)} samples")
            print(f"    Sample: {ds[0]}")
            print()
        except Exception as e:
            print(f"  {name}: ERROR - {e}")

if __name__ == "__main__":
    test_load()