src/verify_dataset.py

#!/usr/bin/env python3
"""Verify SPEED-Bench dataset structure (Claim 1)."""
import json
from pathlib import Path
from datasets import load_from_disk

def verify_qualitative():
    ds = load_from_disk('artifacts/dataset/qualitative')['test']
    print(f"Qualitative split: {len(ds)} samples")
    print(f"Columns: {list(ds.features.keys())}")
    print(f"First sample: {json.dumps(ds[0], indent=2, default=str)[:800]}")
    
    # Check categories
    if 'category' in ds.features:
        categories = {}
        for sample in ds:
            cat = sample['category']
            categories[cat] = categories.get(cat, 0) + 1
        print(f"\nCategories ({len(categories)} total):")
        for cat, count in sorted(categories.items()):
            print(f"  {cat}: {count}")
        
        # The paper says 11 categories with 80 samples each
        expected_cats = 11
        expected_per_cat = 80
        print(f"\nExpected: {expected_cats} categories × {expected_per_cat} = {expected_cats * expected_per_cat}")
        print(f"Actual: {len(categories)} categories, total {sum(categories.values())}")
        
        all_80 = all(count == expected_per_cat for count in categories.values())
        print(f"All categories have exactly {expected_per_cat} samples: {all_80}")
        
        # Check metadata fields
        expected_fields = ['category', 'subcategory', 'multiturn', 'difficulty']
        available_fields = [f for f in expected_fields if f in ds.features]
        print(f"\nMetadata fields present: {available_fields}")
        
        return len(categories) == expected_cats and all_80
    return False

def verify_throughput():
    isls = ['1k', '2k', '8k', '16k', '32k']
    entropy_levels = ['high_entropy', 'mixed', 'low_entropy']
    
    total_samples = 0
    for isl in isls:
        ds = load_from_disk(f'artifacts/dataset/throughput_{isl}')['test']
        print(f"\nThroughput {isl}: {len(ds)} samples")
        
        # Check entropy categories
        if 'entropy' in ds.features:
            entropies = {}
            for sample in ds:
                e = sample['entropy']
                entropies[e] = entropies.get(e, 0) + 1
            print(f"  Entropy distribution: {entropies}")
        
        total_samples += len(ds)
    
    expected_total = 5 * 1536
    print(f"\nTotal throughput samples: {total_samples} (expected {expected_total})")
    return total_samples == expected_total

def main():
    print("=" * 60)
    print("VERIFYING SPEED-BENCH DATASET STRUCTURE (CLAIM 1)")
    print("=" * 60)
    
    q_ok = verify_qualitative()
    t_ok = verify_throughput()
    
    print("\n" + "=" * 60)
    print(f"Claim 1 VERIFIED: {q_ok and t_ok}")
    print("=" * 60)

if __name__ == "__main__":
    main()