src/verify_diversity.py

#!/usr/bin/env python3
"""Toy verification of Claim 2: Greedy selection algorithm produces diverse subsets."""
import json
import numpy as np
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity

def greedy_selection_with_swap(embeddings, k, max_iter=50, seed=None):
    """
    Algorithm 1: Greedy Selection with Local Swap Refinement.
    embeddings: (N, d) array, already normalized
    k: target subset size
    """
    N = len(embeddings)
    if seed is not None:
        np.random.seed(seed)
    
    # Step 1: Greedy selection
    S = [np.random.randint(N)]
    m = embeddings[S[0]].copy()
    
    while len(S) < k:
        scores = embeddings @ m
        mask = np.ones(N, dtype=bool)
        mask[S] = False
        i_star = np.argmin(scores[mask])
        candidates = np.where(mask)[0]
        i_star = candidates[i_star]
        S.append(i_star)
        m += embeddings[i_star]
    
    # Step 2: Local swap refinement
    S_set = set(S)
    not_S = list(set(range(N)) - S_set)
    
    for iteration in range(max_iter):
        improved = False
        for i_o in list(S_set):
            for i_i in not_S:
                delta = 0
                for j in S_set:
                    if j != i_o:
                        delta += (embeddings[i_i] @ embeddings[j] - embeddings[i_o] @ embeddings[j])
                if delta < -1e-9:
                    S_set.remove(i_o)
                    S_set.add(i_i)
                    not_S.remove(i_i)
                    not_S.append(i_o)
                    improved = True
                    break
            if improved:
                break
        if not improved:
            break
    
    return list(S_set)

def avg_pairwise_similarity(embeddings):
    """Compute average pairwise cosine similarity (excluding self)."""
    sim_matrix = cosine_similarity(embeddings)
    np.fill_diagonal(sim_matrix, np.nan)
    return np.nanmean(sim_matrix)

def main():
    print("=" * 60)
    print("CLAIM 2 TOY VERIFICATION: Greedy Selection Algorithm")
    print("=" * 60)
    
    # Load SPEED-Bench qualitative split as our test pool
    from datasets import load_from_disk
    ds = load_from_disk('artifacts/dataset/qualitative')['test']
    
    texts = [sample['turns'][0] if sample['turns'] else "" for sample in ds]
    print(f"Total pool size: {len(texts)} samples")
    
    # Compute embeddings
    print("Computing embeddings...")
    model = SentenceTransformer('all-MiniLM-L6-v2')
    embeddings = model.encode(texts, show_progress_bar=True, convert_to_numpy=True)
    embeddings = embeddings / (np.linalg.norm(embeddings, axis=1, keepdims=True) + 1e-10)
    print(f"Embeddings shape: {embeddings.shape}")
    
    # Test on each category: treat the full dataset as pool, select 80 per category
    # This is a synthetic test - we're verifying the ALGORITHM, not the exact SPEED-Bench curation
    from collections import defaultdict
    cat_indices = defaultdict(list)
    for i, sample in enumerate(ds):
        cat_indices[sample['category']].append(i)
    
    results = {}
    
    for cat, indices in sorted(cat_indices.items()):
        print(f"\n--- Category: {cat} ({len(indices)} samples in pool) ---")
        cat_embeddings = embeddings[indices]
        k = min(80, len(indices) // 2)  # Select half the pool for toy test
        
        # Greedy + Swap
        greedy_idx = greedy_selection_with_swap(cat_embeddings, k, seed=42)
        greedy_sim = avg_pairwise_similarity(cat_embeddings[greedy_idx])
        print(f"  Greedy+Swap (k={k}): avg sim = {greedy_sim:.4f}")
        
        # Random baseline (multiple seeds)
        rand_sims = []
        for seed in range(5):
            np.random.seed(seed)
            rand_idx = np.random.choice(len(cat_embeddings), size=k, replace=False)
            rand_sim = avg_pairwise_similarity(cat_embeddings[rand_idx])
            rand_sims.append(rand_sim)
        rand_avg = np.mean(rand_sims)
        print(f"  Random (k={k}): avg sim = {rand_avg:.4f} (std={np.std(rand_sims):.4f})")
        
        # Full pool (for reference)
        full_sim = avg_pairwise_similarity(cat_embeddings)
        print(f"  Full pool: avg sim = {full_sim:.4f}")
        
        results[cat] = {
            'greedy_sim': float(greedy_sim),
            'random_sim': float(rand_avg),
            'full_sim': float(full_sim),
            'greedy_better': bool(greedy_sim < rand_avg),
            'k': int(k)
        }
    
    # Overall summary
    print("\n" + "=" * 60)
    print("SUMMARY")
    print("=" * 60)
    
    greedy_vals = [v['greedy_sim'] for v in results.values()]
    rand_vals = [v['random_sim'] for v in results.values()]
    
    print(f"Avg pairwise similarity across categories:")
    print(f"  Greedy+Swap: {np.mean(greedy_vals):.4f}")
    print(f"  Random:      {np.mean(rand_vals):.4f}")
    print(f"  Improvement: {np.mean(rand_vals) - np.mean(greedy_vals):.4f}")
    
    n_better = sum(1 for v in results.values() if v['greedy_better'])
    print(f"\nGreedy+Swap better than random in {n_better}/{len(results)} categories")
    
    claim2_verified = n_better >= len(results) * 0.7  # At least 70%
    print(f"Claim 2 (toy) VERIFIED: {claim2_verified}")
    
    # Save results
    with open('outputs/claim2_diversity.json', 'w') as f:
        json.dump({
            'overall': {
                'greedy_avg': float(np.mean(greedy_vals)),
                'random_avg': float(np.mean(rand_vals)),
                'improvement': float(np.mean(rand_vals) - np.mean(greedy_vals)),
                'n_better': int(n_better),
                'n_total': len(results),
                'verified': claim2_verified
            },
            'per_category': results
        }, f, indent=2)
    print("\nResults saved to outputs/claim2_diversity.json")

if __name__ == "__main__":
    main()