wisent 0.7.901__py3-none-any.whl → 0.7.1045__py3-none-any.whl

wisent/examples/scripts/threshold_analysis.py

@@ -1,434 +0,0 @@
-"""
-Threshold Analysis for RepScan.
-
-Analyzes sensitivity of diagnosis to threshold choices:
-- ROC curves for existence threshold
-- Precision/recall tradeoff
-- Null distribution analysis
-- Synthetic validation
-
-Usage:
-    python -m wisent.examples.scripts.threshold_analysis --model Qwen/Qwen3-8B
-"""
-
-import argparse
-import json
-import subprocess
-from pathlib import Path
-from typing import Dict, List, Any, Optional, Tuple
-from dataclasses import dataclass, field, asdict
-import random
-
-import torch
-import numpy as np
-from sklearn.metrics import roc_curve, auc, precision_recall_curve
-
-S3_BUCKET = "wisent-bucket"
-S3_PREFIX = "threshold_analysis"
-
-
-def s3_upload_file(local_path: Path, model_name: str) -> None:
-    """Upload a single file to S3."""
-    model_prefix = model_name.replace('/', '_')
-    s3_path = f"s3://{S3_BUCKET}/{S3_PREFIX}/{model_prefix}/{local_path.name}"
-    try:
-        subprocess.run(
-            ["aws", "s3", "cp", str(local_path), s3_path, "--quiet"],
-            check=True,
-            capture_output=True,
-        )
-        print(f"  Uploaded to S3: {s3_path}")
-    except Exception as e:
-        print(f"  S3 upload failed: {e}")
-
-
-@dataclass
-class ThresholdAnalysisResult:
-    """Result of threshold analysis."""
-    # Existence threshold analysis
-    existence_thresholds: List[float]
-    existence_tpr: List[float]  # True positive rate
-    existence_fpr: List[float]  # False positive rate
-    existence_auc: float
-    optimal_existence_threshold: float
-    
-    # Gap threshold analysis
-    gap_thresholds: List[float]
-    gap_precision: List[float]
-    gap_recall: List[float]
-    gap_f1: List[float]
-    optimal_gap_threshold: float
-    
-    # Null distribution stats
-    null_mean_knn: float
-    null_std_knn: float
-    null_mean_linear: float
-    null_std_linear: float
-    
-    # Sensitivity analysis
-    sensitivity_matrix: Dict[str, Dict[str, float]]  # threshold -> diagnosis distribution
-
-
-def generate_null_distribution(
-    model: "WisentModel",
-    n_samples: int = 100,
-    hidden_dim: int = 4096,
-) -> Tuple[List[float], List[float]]:
-    """
-    Generate null distribution by testing random/nonsense data.
-    
-    Args:
-        model: WisentModel instance
-        n_samples: Number of random samples
-        hidden_dim: Hidden dimension
-        
-    Returns:
-        (knn_scores, linear_scores) for random data
-    """
-    from wisent.core.geometry_runner import compute_knn_accuracy, compute_linear_probe_accuracy
-    
-    knn_scores = []
-    linear_scores = []
-    
-    for _ in range(n_samples):
-        # Generate random activations (no real signal)
-        pos = torch.randn(50, hidden_dim)
-        neg = torch.randn(50, hidden_dim)
-        
-        knn = compute_knn_accuracy(pos, neg, k=10)
-        linear = compute_linear_probe_accuracy(pos, neg)
-        
-        knn_scores.append(knn)
-        linear_scores.append(linear)
-    
-    return knn_scores, linear_scores
-
-
-def generate_synthetic_data(
-    structure: str,
-    n_samples: int = 50,
-    hidden_dim: int = 100,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    """
-    Generate synthetic data with known structure for validation.
-    
-    Args:
-        structure: 'linear', 'xor', 'spirals', 'random'
-        n_samples: Samples per class
-        hidden_dim: Dimension
-        
-    Returns:
-        (pos_activations, neg_activations)
-    """
-    if structure == "linear":
-        # Linear separable: positive class shifted in one direction
-        direction = torch.randn(hidden_dim)
-        direction = direction / direction.norm()
-        
-        pos = torch.randn(n_samples, hidden_dim) + 2 * direction
-        neg = torch.randn(n_samples, hidden_dim) - 2 * direction
-        
-    elif structure == "xor":
-        # XOR pattern: nonlinear but separable
-        base = torch.randn(n_samples, hidden_dim)
-        
-        # Positive: (high dim1 AND high dim2) OR (low dim1 AND low dim2)
-        pos_mask1 = (base[:n_samples//2, 0] > 0) & (base[:n_samples//2, 1] > 0)
-        pos_mask2 = (base[n_samples//2:, 0] < 0) & (base[n_samples//2:, 1] < 0)
-        
-        pos = torch.randn(n_samples, hidden_dim)
-        pos[:n_samples//2, 0] = torch.abs(pos[:n_samples//2, 0]) + 1
-        pos[:n_samples//2, 1] = torch.abs(pos[:n_samples//2, 1]) + 1
-        pos[n_samples//2:, 0] = -torch.abs(pos[n_samples//2:, 0]) - 1
-        pos[n_samples//2:, 1] = -torch.abs(pos[n_samples//2:, 1]) - 1
-        
-        neg = torch.randn(n_samples, hidden_dim)
-        neg[:n_samples//2, 0] = torch.abs(neg[:n_samples//2, 0]) + 1
-        neg[:n_samples//2, 1] = -torch.abs(neg[:n_samples//2, 1]) - 1
-        neg[n_samples//2:, 0] = -torch.abs(neg[n_samples//2:, 0]) - 1
-        neg[n_samples//2:, 1] = torch.abs(neg[n_samples//2:, 1]) + 1
-        
-    elif structure == "spirals":
-        # Interleaved spirals: nonlinear separable
-        t_pos = torch.linspace(0, 4*np.pi, n_samples)
-        t_neg = torch.linspace(0, 4*np.pi, n_samples) + np.pi
-        
-        pos = torch.zeros(n_samples, hidden_dim)
-        pos[:, 0] = t_pos * torch.cos(t_pos) + 0.5 * torch.randn(n_samples)
-        pos[:, 1] = t_pos * torch.sin(t_pos) + 0.5 * torch.randn(n_samples)
-        pos[:, 2:] = torch.randn(n_samples, hidden_dim - 2) * 0.1
-        
-        neg = torch.zeros(n_samples, hidden_dim)
-        neg[:, 0] = t_neg * torch.cos(t_neg) + 0.5 * torch.randn(n_samples)
-        neg[:, 1] = t_neg * torch.sin(t_neg) + 0.5 * torch.randn(n_samples)
-        neg[:, 2:] = torch.randn(n_samples, hidden_dim - 2) * 0.1
-        
-    else:  # random
-        pos = torch.randn(n_samples, hidden_dim)
-        neg = torch.randn(n_samples, hidden_dim)
-    
-    return pos, neg
-
-
-def compute_roc_for_existence(
-    real_results: List[Dict],
-    null_scores: List[float],
-) -> Tuple[List[float], List[float], List[float], float]:
-    """
-    Compute ROC curve for existence threshold.
-    
-    Args:
-        real_results: Results from real benchmarks
-        null_scores: kNN scores from null distribution
-        
-    Returns:
-        (thresholds, tpr, fpr, auc)
-    """
-    # Labels: 1 for real data (should be detected), 0 for null (should not)
-    real_knn = [r["nonlinear_metrics"]["knn_accuracy_k10"] for r in real_results]
-    
-    scores = real_knn + null_scores
-    labels = [1] * len(real_knn) + [0] * len(null_scores)
-    
-    fpr, tpr, thresholds = roc_curve(labels, scores)
-    roc_auc = auc(fpr, tpr)
-    
-    return thresholds.tolist(), tpr.tolist(), fpr.tolist(), roc_auc
-
-
-def compute_precision_recall_for_gap(
-    results: List[Dict],
-    ground_truth_linear: List[bool],
-) -> Tuple[List[float], List[float], List[float], List[float]]:
-    """
-    Compute precision-recall for gap threshold (linear vs nonlinear).
-    
-    Args:
-        results: Results from benchmarks
-        ground_truth_linear: Ground truth labels (True = linear, False = nonlinear)
-        
-    Returns:
-        (thresholds, precision, recall, f1)
-    """
-    # Gap = signal_strength - linear_probe_accuracy
-    gaps = [r["signal_strength"] - r["linear_probe_accuracy"] for r in results]
-    
-    # Labels: 1 for nonlinear (gap > threshold), 0 for linear
-    labels = [0 if gt else 1 for gt in ground_truth_linear]
-    
-    precision, recall, thresholds = precision_recall_curve(labels, gaps)
-    
-    # Compute F1
-    f1 = [2 * p * r / (p + r + 1e-10) for p, r in zip(precision, recall)]
-    
-    return thresholds.tolist(), precision.tolist(), recall.tolist(), f1
-
-
-def run_sensitivity_analysis(
-    results: List[Dict],
-    existence_thresholds: List[float] = [0.5, 0.55, 0.6, 0.65, 0.7],
-    gap_thresholds: List[float] = [0.05, 0.10, 0.15, 0.20, 0.25],
-) -> Dict[str, Dict[str, float]]:
-    """
-    Run sensitivity analysis across threshold combinations.
-    
-    Args:
-        results: Results from benchmarks
-        existence_thresholds: Thresholds to test for existence
-        gap_thresholds: Thresholds to test for gap
-        
-    Returns:
-        Nested dict: {exist_thresh: {gap_thresh: {diagnosis: percentage}}}
-    """
-    sensitivity = {}
-    
-    for exist_t in existence_thresholds:
-        sensitivity[str(exist_t)] = {}
-        
-        for gap_t in gap_thresholds:
-            diagnoses = {"LINEAR": 0, "NONLINEAR": 0, "NO_SIGNAL": 0}
-            
-            for r in results:
-                signal = r["signal_strength"]
-                gap = signal - r["linear_probe_accuracy"]
-                
-                if signal < exist_t:
-                    diagnoses["NO_SIGNAL"] += 1
-                elif gap < gap_t:
-                    diagnoses["LINEAR"] += 1
-                else:
-                    diagnoses["NONLINEAR"] += 1
-            
-            total = len(results)
-            sensitivity[str(exist_t)][str(gap_t)] = {
-                k: v / total * 100 for k, v in diagnoses.items()
-            }
-    
-    return sensitivity
-
-
-def load_diagnosis_results(model_name: str, output_dir: Path) -> List[Dict]:
-    """Load all diagnosis results."""
-    model_prefix = model_name.replace('/', '_')
-    
-    # Try to download from S3
-    try:
-        subprocess.run(
-            ["aws", "s3", "sync", 
-             f"s3://{S3_BUCKET}/direction_discovery/{model_prefix}/",
-             str(output_dir / "diagnosis"),
-             "--quiet"],
-            check=False,
-            capture_output=True,
-        )
-    except Exception:
-        pass
-    
-    # Load all results
-    all_results = []
-    diagnosis_dir = output_dir / "diagnosis"
-    
-    if diagnosis_dir.exists():
-        for f in diagnosis_dir.glob(f"{model_prefix}_*.json"):
-            if "summary" not in f.name:
-                with open(f) as fp:
-                    data = json.load(fp)
-                    all_results.extend(data.get("results", []))
-    
-    return all_results
-
-
-def run_threshold_analysis(model_name: str):
-    """
-    Run full threshold analysis.
-    
-    Args:
-        model_name: Model to analyze
-    """
-    print("=" * 70)
-    print("THRESHOLD ANALYSIS")
-    print("=" * 70)
-    print(f"Model: {model_name}")
-    
-    output_dir = Path("/tmp/threshold_analysis")
-    output_dir.mkdir(parents=True, exist_ok=True)
-    
-    # Load diagnosis results
-    results = load_diagnosis_results(model_name, output_dir)
-    if not results:
-        print("ERROR: No diagnosis results found.")
-        return
-    
-    print(f"Loaded {len(results)} results")
-    
-    # 1. Generate null distribution
-    print("\n1. Generating null distribution...")
-    null_knn, null_linear = generate_null_distribution(None, n_samples=100, hidden_dim=4096)
-    
-    print(f"   Null kNN: mean={np.mean(null_knn):.3f}, std={np.std(null_knn):.3f}")
-    print(f"   Null linear: mean={np.mean(null_linear):.3f}, std={np.std(null_linear):.3f}")
-    
-    # 2. ROC for existence threshold
-    print("\n2. Computing ROC for existence threshold...")
-    thresholds, tpr, fpr, roc_auc = compute_roc_for_existence(results, null_knn)
-    
-    # Find optimal threshold (Youden's J)
-    j_scores = [t - f for t, f in zip(tpr, fpr)]
-    optimal_idx = np.argmax(j_scores)
-    optimal_exist = thresholds[optimal_idx] if optimal_idx < len(thresholds) else 0.6
-    
-    print(f"   AUC: {roc_auc:.3f}")
-    print(f"   Optimal existence threshold: {optimal_exist:.3f}")
-    
-    # 3. Synthetic validation
-    print("\n3. Synthetic validation...")
-    from wisent.core.geometry_runner import compute_knn_accuracy, compute_linear_probe_accuracy
-    
-    synthetic_results = {}
-    for structure in ["linear", "xor", "spirals", "random"]:
-        pos, neg = generate_synthetic_data(structure)
-        knn = compute_knn_accuracy(pos, neg, k=10)
-        linear = compute_linear_probe_accuracy(pos, neg)
-        gap = knn - linear
-        
-        synthetic_results[structure] = {
-            "knn": knn,
-            "linear": linear,
-            "gap": gap,
-        }
-        print(f"   {structure}: kNN={knn:.3f}, linear={linear:.3f}, gap={gap:.3f}")
-    
-    # Validate that gap threshold separates linear from nonlinear
-    linear_gap = synthetic_results["linear"]["gap"]
-    xor_gap = synthetic_results["xor"]["gap"]
-    spirals_gap = synthetic_results["spirals"]["gap"]
-    
-    # Good gap threshold should be > linear_gap and < min(xor_gap, spirals_gap)
-    optimal_gap = (linear_gap + min(xor_gap, spirals_gap)) / 2
-    print(f"\n   Suggested gap threshold: {optimal_gap:.3f}")
-    
-    # 4. Sensitivity analysis
-    print("\n4. Running sensitivity analysis...")
-    sensitivity = run_sensitivity_analysis(results)
-    
-    print("\n   Diagnosis distribution (% of benchmarks):")
-    print("   " + "-" * 60)
-    print(f"   {'Exist':>6} | {'Gap':>6} | {'LINEAR':>8} | {'NONLINEAR':>10} | {'NO_SIGNAL':>10}")
-    print("   " + "-" * 60)
-    
-    for exist_t, gap_data in sensitivity.items():
-        for gap_t, diagnoses in gap_data.items():
-            print(f"   {exist_t:>6} | {gap_t:>6} | {diagnoses['LINEAR']:>7.1f}% | "
-                  f"{diagnoses['NONLINEAR']:>9.1f}% | {diagnoses['NO_SIGNAL']:>9.1f}%")
-    
-    # 5. Save results
-    analysis_result = ThresholdAnalysisResult(
-        existence_thresholds=thresholds[:100],  # Limit for JSON
-        existence_tpr=tpr[:100],
-        existence_fpr=fpr[:100],
-        existence_auc=roc_auc,
-        optimal_existence_threshold=float(optimal_exist),
-        gap_thresholds=[0.05, 0.10, 0.15, 0.20, 0.25],
-        gap_precision=[],  # Would need ground truth
-        gap_recall=[],
-        gap_f1=[],
-        optimal_gap_threshold=float(optimal_gap),
-        null_mean_knn=float(np.mean(null_knn)),
-        null_std_knn=float(np.std(null_knn)),
-        null_mean_linear=float(np.mean(null_linear)),
-        null_std_linear=float(np.std(null_linear)),
-        sensitivity_matrix=sensitivity,
-    )
-    
-    model_prefix = model_name.replace('/', '_')
-    results_file = output_dir / f"{model_prefix}_threshold_analysis.json"
-    
-    with open(results_file, "w") as f:
-        json.dump(asdict(analysis_result), f, indent=2)
-    
-    print(f"\nResults saved to: {results_file}")
-    s3_upload_file(results_file, model_name)
-    
-    # Summary
-    print("\n" + "=" * 70)
-    print("RECOMMENDATIONS")
-    print("=" * 70)
-    print(f"\n1. Existence threshold: {optimal_exist:.2f}")
-    print(f"   - Based on ROC analysis (AUC={roc_auc:.3f})")
-    print(f"   - Null distribution: kNN={np.mean(null_knn):.3f} ± {np.std(null_knn):.3f}")
-    
-    print(f"\n2. Gap threshold: {optimal_gap:.2f}")
-    print(f"   - Based on synthetic validation")
-    print(f"   - Linear structure gap: {linear_gap:.3f}")
-    print(f"   - XOR structure gap: {xor_gap:.3f}")
-    print(f"   - Spirals structure gap: {spirals_gap:.3f}")
-    
-    return analysis_result
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Threshold analysis for RepScan")
-    parser.add_argument("--model", type=str, default="Qwen/Qwen3-8B", help="Model to analyze")
-    args = parser.parse_args()
-    
-    run_threshold_analysis(args.model)