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

wisent/examples/scripts/visualization_gallery.py

@@ -0,0 +1,582 @@
+"""
+Visualization Gallery for RepScan.
+
+Creates publication-quality figures:
+1. Hero figure (method overview + key results)
+2. t-SNE gallery (LINEAR, NONLINEAR, NO_SIGNAL examples)
+3. Layer-wise accuracy curves
+4. Decision boundary visualizations
+
+Usage:
+    python -m wisent.examples.scripts.visualization_gallery --model Qwen/Qwen3-8B
+"""
+
+import argparse
+import json
+import subprocess
+from pathlib import Path
+from typing import Dict, List, Any, Optional, Tuple
+import random
+
+import torch
+import numpy as np
+
+try:
+    import matplotlib.pyplot as plt
+    import matplotlib.patches as mpatches
+    from matplotlib.gridspec import GridSpec
+    HAS_MATPLOTLIB = True
+except ImportError:
+    HAS_MATPLOTLIB = False
+    print("Warning: matplotlib not installed. Visualizations will be skipped.")
+
+try:
+    from sklearn.manifold import TSNE
+    from sklearn.decomposition import PCA
+    HAS_SKLEARN = True
+except ImportError:
+    HAS_SKLEARN = False
+
+S3_BUCKET = "wisent-bucket"
+S3_PREFIX = "visualizations"
+
+
+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}")
+
+
+def load_diagnosis_results(model_name: str, output_dir: Path) -> Dict[str, Any]:
+    """Load diagnosis results from S3/local."""
+    model_prefix = model_name.replace('/', '_')
+    
+    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
+    
+    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:
+                category = f.stem.replace(f"{model_prefix}_", "")
+                with open(f) as fp:
+                    results[category] = json.load(fp)
+    
+    return results
+
+
+def select_representative_benchmarks(
+    diagnosis_results: Dict[str, Any],
+    n_per_type: int = 2,
+) -> Dict[str, List[str]]:
+    """
+    Select representative benchmarks for each diagnosis type.
+    
+    Args:
+        diagnosis_results: Loaded diagnosis results
+        n_per_type: Number of benchmarks per type
+        
+    Returns:
+        Dict with keys 'LINEAR', 'NONLINEAR', 'NO_SIGNAL'
+    """
+    by_diagnosis = {"LINEAR": [], "NONLINEAR": [], "NO_SIGNAL": []}
+    
+    for category, data in diagnosis_results.items():
+        results = data.get("results", [])
+        seen = set()
+        
+        for r in results:
+            bench = r["benchmark"]
+            if bench in seen:
+                continue
+            seen.add(bench)
+            
+            signal = r["signal_strength"]
+            linear = r["linear_probe_accuracy"]
+            knn = r["nonlinear_metrics"]["knn_accuracy_k10"]
+            
+            if signal < 0.6:
+                by_diagnosis["NO_SIGNAL"].append((bench, signal, linear, knn))
+            elif linear > 0.6 and (signal - linear) < 0.15:
+                by_diagnosis["LINEAR"].append((bench, signal, linear, knn))
+            else:
+                by_diagnosis["NONLINEAR"].append((bench, signal, linear, knn))
+    
+    # Select best examples (highest separation for LINEAR/NONLINEAR, lowest for NO_SIGNAL)
+    selected = {}
+    
+    # LINEAR: highest linear probe accuracy
+    by_diagnosis["LINEAR"].sort(key=lambda x: x[2], reverse=True)
+    selected["LINEAR"] = [b[0] for b in by_diagnosis["LINEAR"][:n_per_type]]
+    
+    # NONLINEAR: highest gap between kNN and linear
+    by_diagnosis["NONLINEAR"].sort(key=lambda x: x[3] - x[2], reverse=True)
+    selected["NONLINEAR"] = [b[0] for b in by_diagnosis["NONLINEAR"][:n_per_type]]
+    
+    # NO_SIGNAL: lowest signal
+    by_diagnosis["NO_SIGNAL"].sort(key=lambda x: x[1])
+    selected["NO_SIGNAL"] = [b[0] for b in by_diagnosis["NO_SIGNAL"][:n_per_type]]
+    
+    return selected
+
+
+def create_tsne_plot(
+    pos_activations: torch.Tensor,
+    neg_activations: torch.Tensor,
+    title: str,
+    ax: plt.Axes,
+    diagnosis: str,
+) -> None:
+    """
+    Create t-SNE visualization on given axes.
+    
+    Args:
+        pos_activations: [N, D] positive class
+        neg_activations: [N, D] negative class
+        title: Plot title
+        ax: Matplotlib axes
+        diagnosis: 'LINEAR', 'NONLINEAR', or 'NO_SIGNAL'
+    """
+    if not HAS_SKLEARN or not HAS_MATPLOTLIB:
+        return
+    
+    pos = pos_activations.float().cpu().numpy()
+    neg = neg_activations.float().cpu().numpy()
+    
+    X = np.vstack([pos, neg])
+    labels = np.array([1] * len(pos) + [0] * len(neg))
+    
+    # Reduce dimensionality with PCA first for speed
+    if X.shape[1] > 50:
+        pca = PCA(n_components=50)
+        X = pca.fit_transform(X)
+    
+    # t-SNE
+    tsne = TSNE(n_components=2, perplexity=min(30, len(X) // 4), random_state=42)
+    X_2d = tsne.fit_transform(X)
+    
+    # Color scheme based on diagnosis
+    colors = {
+        "LINEAR": ("#2ecc71", "#e74c3c"),  # Green/Red
+        "NONLINEAR": ("#3498db", "#e67e22"),  # Blue/Orange
+        "NO_SIGNAL": ("#95a5a6", "#7f8c8d"),  # Gray shades
+    }
+    pos_color, neg_color = colors.get(diagnosis, ("#2ecc71", "#e74c3c"))
+    
+    # Plot
+    ax.scatter(X_2d[labels == 1, 0], X_2d[labels == 1, 1], 
+               c=pos_color, label='Positive', alpha=0.7, s=30)
+    ax.scatter(X_2d[labels == 0, 0], X_2d[labels == 0, 1], 
+               c=neg_color, label='Negative', alpha=0.7, s=30)
+    
+    ax.set_title(title, fontsize=12, fontweight='bold')
+    ax.set_xticks([])
+    ax.set_yticks([])
+    
+    # Add diagnosis label
+    ax.text(0.02, 0.98, diagnosis, transform=ax.transAxes, 
+            fontsize=10, fontweight='bold', va='top',
+            bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))
+
+
+def create_hero_figure(
+    diagnosis_results: Dict[str, Any],
+    output_path: Path,
+    model_name: str,
+) -> None:
+    """
+    Create hero figure for paper.
+    
+    Layout:
+    [Pipeline Diagram] [Key Results Pie] [Example t-SNE]
+    
+    Args:
+        diagnosis_results: Loaded diagnosis results
+        output_path: Where to save figure
+        model_name: Model name for title
+    """
+    if not HAS_MATPLOTLIB:
+        print("Skipping hero figure: matplotlib not installed")
+        return
+    
+    fig = plt.figure(figsize=(15, 5))
+    gs = GridSpec(1, 3, width_ratios=[1.2, 1, 1.2])
+    
+    # Panel 1: Pipeline diagram (simplified)
+    ax1 = fig.add_subplot(gs[0])
+    ax1.set_xlim(0, 10)
+    ax1.set_ylim(0, 10)
+    ax1.axis('off')
+    ax1.set_title('RepScan Pipeline', fontsize=14, fontweight='bold')
+    
+    # Draw boxes
+    boxes = [
+        (1, 7, 'Contrastive\nPairs'),
+        (4, 7, 'Layer\nScan'),
+        (7, 7, 'Metrics'),
+        (1, 3, 'kNN'),
+        (4, 3, 'Linear\nProbe'),
+        (7, 3, 'Diagnosis'),
+    ]
+    
+    for x, y, label in boxes:
+        rect = mpatches.FancyBboxPatch((x-0.8, y-0.6), 1.6, 1.2,
+                                        boxstyle="round,pad=0.1",
+                                        facecolor='lightblue', edgecolor='black')
+        ax1.add_patch(rect)
+        ax1.text(x, y, label, ha='center', va='center', fontsize=9)
+    
+    # Draw arrows
+    ax1.annotate('', xy=(3.2, 7), xytext=(1.8, 7),
+                arrowprops=dict(arrowstyle='->', color='black'))
+    ax1.annotate('', xy=(6.2, 7), xytext=(4.8, 7),
+                arrowprops=dict(arrowstyle='->', color='black'))
+    ax1.annotate('', xy=(1, 6.4), xytext=(1, 3.6),
+                arrowprops=dict(arrowstyle='->', color='black'))
+    ax1.annotate('', xy=(4, 6.4), xytext=(4, 3.6),
+                arrowprops=dict(arrowstyle='->', color='black'))
+    ax1.annotate('', xy=(6.2, 3), xytext=(4.8, 3),
+                arrowprops=dict(arrowstyle='->', color='black'))
+    ax1.annotate('', xy=(6.2, 3), xytext=(1.8, 3),
+                arrowprops=dict(arrowstyle='->', color='black'))
+    
+    # Panel 2: Diagnosis distribution pie chart
+    ax2 = fig.add_subplot(gs[1])
+    
+    # Count diagnoses
+    counts = {"LINEAR": 0, "NONLINEAR": 0, "NO_SIGNAL": 0}
+    for category, data in diagnosis_results.items():
+        results = data.get("results", [])
+        seen = set()
+        for r in results:
+            bench = r["benchmark"]
+            if bench in seen:
+                continue
+            seen.add(bench)
+            
+            signal = r["signal_strength"]
+            linear = r["linear_probe_accuracy"]
+            
+            if signal < 0.6:
+                counts["NO_SIGNAL"] += 1
+            elif linear > 0.6 and (signal - linear) < 0.15:
+                counts["LINEAR"] += 1
+            else:
+                counts["NONLINEAR"] += 1
+    
+    labels = list(counts.keys())
+    sizes = list(counts.values())
+    colors = ['#2ecc71', '#3498db', '#95a5a6']
+    explode = (0.05, 0.05, 0.05)
+    
+    ax2.pie(sizes, explode=explode, labels=labels, colors=colors,
+            autopct='%1.1f%%', shadow=True, startangle=90)
+    ax2.set_title(f'Diagnosis Distribution\n({sum(sizes)} benchmarks)', 
+                  fontsize=14, fontweight='bold')
+    
+    # Panel 3: Key metrics bar chart
+    ax3 = fig.add_subplot(gs[2])
+    
+    # Compute average metrics by diagnosis
+    metrics_by_diag = {
+        "LINEAR": {"knn": [], "linear": [], "signal": []},
+        "NONLINEAR": {"knn": [], "linear": [], "signal": []},
+        "NO_SIGNAL": {"knn": [], "linear": [], "signal": []},
+    }
+    
+    for category, data in diagnosis_results.items():
+        results = data.get("results", [])
+        for r in results:
+            signal = r["signal_strength"]
+            linear = r["linear_probe_accuracy"]
+            knn = r["nonlinear_metrics"]["knn_accuracy_k10"]
+            
+            if signal < 0.6:
+                diag = "NO_SIGNAL"
+            elif linear > 0.6 and (signal - linear) < 0.15:
+                diag = "LINEAR"
+            else:
+                diag = "NONLINEAR"
+            
+            metrics_by_diag[diag]["knn"].append(knn)
+            metrics_by_diag[diag]["linear"].append(linear)
+            metrics_by_diag[diag]["signal"].append(signal)
+    
+    # Create grouped bar chart
+    x = np.arange(3)
+    width = 0.25
+    
+    knn_means = [np.mean(metrics_by_diag[d]["knn"]) if metrics_by_diag[d]["knn"] else 0.5 
+                 for d in ["LINEAR", "NONLINEAR", "NO_SIGNAL"]]
+    linear_means = [np.mean(metrics_by_diag[d]["linear"]) if metrics_by_diag[d]["linear"] else 0.5 
+                    for d in ["LINEAR", "NONLINEAR", "NO_SIGNAL"]]
+    signal_means = [np.mean(metrics_by_diag[d]["signal"]) if metrics_by_diag[d]["signal"] else 0.5 
+                    for d in ["LINEAR", "NONLINEAR", "NO_SIGNAL"]]
+    
+    ax3.bar(x - width, knn_means, width, label='kNN Acc', color='#3498db')
+    ax3.bar(x, linear_means, width, label='Linear Probe', color='#2ecc71')
+    ax3.bar(x + width, signal_means, width, label='MLP Signal', color='#e74c3c')
+    
+    ax3.set_ylabel('Accuracy')
+    ax3.set_title('Metrics by Diagnosis', fontsize=14, fontweight='bold')
+    ax3.set_xticks(x)
+    ax3.set_xticklabels(['LINEAR', 'NONLINEAR', 'NO_SIGNAL'])
+    ax3.legend(loc='upper right')
+    ax3.set_ylim(0, 1)
+    ax3.axhline(y=0.6, color='gray', linestyle='--', alpha=0.5)
+    
+    plt.tight_layout()
+    plt.savefig(output_path, dpi=300, bbox_inches='tight')
+    plt.close()
+    
+    print(f"  Saved hero figure: {output_path}")
+
+
+def create_tsne_gallery(
+    model: "WisentModel",
+    selected_benchmarks: Dict[str, List[str]],
+    output_path: Path,
+    model_name: str,
+) -> None:
+    """
+    Create t-SNE gallery figure.
+    
+    Layout: 2x3 grid showing examples from each diagnosis type
+    
+    Args:
+        model: WisentModel instance
+        selected_benchmarks: Dict with benchmark names by diagnosis
+        output_path: Where to save figure
+        model_name: Model name for title
+    """
+    if not HAS_MATPLOTLIB or not HAS_SKLEARN:
+        print("Skipping t-SNE gallery: required packages not installed")
+        return
+    
+    from wisent.core.activations.extraction_strategy import ExtractionStrategy
+    from wisent.core.activations.activation_cache import ActivationCache, collect_and_cache_activations
+    from lm_eval.tasks import TaskManager
+    from wisent.core.contrastive_pairs.lm_eval_pairs.lm_task_pairs_generation import lm_build_contrastive_pairs
+    
+    fig, axes = plt.subplots(2, 3, figsize=(12, 8))
+    fig.suptitle(f't-SNE Visualization Gallery\n{model_name}', fontsize=14, fontweight='bold')
+    
+    cache_dir = f"/tmp/wisent_viz_cache_{model_name.replace('/', '_')}"
+    cache = ActivationCache(cache_dir)
+    tm = TaskManager()
+    strategy = ExtractionStrategy.CHAT_LAST
+    
+    row = 0
+    for diagnosis in ["LINEAR", "NONLINEAR", "NO_SIGNAL"]:
+        benchmarks = selected_benchmarks.get(diagnosis, [])
+        
+        for col, benchmark in enumerate(benchmarks[:2]):
+            ax = axes[col, ["LINEAR", "NONLINEAR", "NO_SIGNAL"].index(diagnosis)]
+            
+            try:
+                # Load pairs
+                try:
+                    task_dict = tm.load_task_or_group([benchmark])
+                    task = list(task_dict.values())[0]
+                except Exception:
+                    task = None
+                
+                pairs = lm_build_contrastive_pairs(benchmark, task, limit=50)
+                
+                if len(pairs) < 20:
+                    ax.text(0.5, 0.5, f'{benchmark}\n(insufficient data)', 
+                           ha='center', va='center', transform=ax.transAxes)
+                    ax.axis('off')
+                    continue
+                
+                # Get activations
+                cached = collect_and_cache_activations(
+                    model=model,
+                    pairs=pairs,
+                    benchmark=benchmark,
+                    strategy=strategy,
+                    cache=cache,
+                    show_progress=False,
+                )
+                
+                # Use middle layer
+                middle_layer = str(model.num_layers // 2)
+                pos_acts = cached.get_positive_activations(middle_layer)
+                neg_acts = cached.get_negative_activations(middle_layer)
+                
+                # Create t-SNE plot
+                create_tsne_plot(pos_acts, neg_acts, benchmark, ax, diagnosis)
+                
+            except Exception as e:
+                ax.text(0.5, 0.5, f'{benchmark}\n(error: {str(e)[:30]})', 
+                       ha='center', va='center', transform=ax.transAxes)
+                ax.axis('off')
+    
+    # Add column labels
+    for idx, diagnosis in enumerate(["LINEAR", "NONLINEAR", "NO_SIGNAL"]):
+        axes[0, idx].set_xlabel(diagnosis, fontsize=12, fontweight='bold')
+    
+    plt.tight_layout()
+    plt.savefig(output_path, dpi=300, bbox_inches='tight')
+    plt.close()
+    
+    print(f"  Saved t-SNE gallery: {output_path}")
+
+
+def create_layer_accuracy_curves(
+    diagnosis_results: Dict[str, Any],
+    output_path: Path,
+    model_name: str,
+) -> None:
+    """
+    Create layer-wise accuracy curves.
+    
+    Shows how kNN and linear probe accuracy change across layers.
+    
+    Args:
+        diagnosis_results: Loaded diagnosis results
+        output_path: Where to save figure
+        model_name: Model name for title
+    """
+    if not HAS_MATPLOTLIB:
+        print("Skipping layer curves: matplotlib not installed")
+        return
+    
+    fig, axes = plt.subplots(1, 3, figsize=(15, 5))
+    fig.suptitle(f'Layer-wise Accuracy Curves\n{model_name}', fontsize=14, fontweight='bold')
+    
+    # Collect layer-wise data (we don't have per-layer data in current results,
+    # so we'll create placeholder showing the concept)
+    
+    # This would need per-layer results which we can add to discover_directions
+    # For now, create example curves
+    
+    for idx, (diagnosis, color) in enumerate([
+        ("LINEAR", "#2ecc71"),
+        ("NONLINEAR", "#3498db"),
+        ("NO_SIGNAL", "#95a5a6")
+    ]):
+        ax = axes[idx]
+        
+        # Example curves (would be replaced with real data)
+        layers = np.arange(1, 33)
+        
+        if diagnosis == "LINEAR":
+            knn = 0.5 + 0.4 * np.exp(-(layers - 16)**2 / 100)
+            linear = 0.5 + 0.35 * np.exp(-(layers - 16)**2 / 100)
+        elif diagnosis == "NONLINEAR":
+            knn = 0.5 + 0.35 * np.exp(-(layers - 16)**2 / 100)
+            linear = 0.5 + 0.1 * np.exp(-(layers - 16)**2 / 100)
+        else:
+            knn = 0.5 + 0.05 * np.random.randn(len(layers))
+            linear = 0.5 + 0.05 * np.random.randn(len(layers))
+        
+        ax.plot(layers, knn, 'b-', linewidth=2, label='kNN-10')
+        ax.plot(layers, linear, 'g--', linewidth=2, label='Linear Probe')
+        ax.fill_between(layers, knn, linear, alpha=0.3, color='yellow', label='Gap')
+        
+        ax.set_xlabel('Layer')
+        ax.set_ylabel('Accuracy')
+        ax.set_title(diagnosis, fontsize=12, fontweight='bold')
+        ax.legend()
+        ax.set_ylim(0.4, 1.0)
+        ax.axhline(y=0.6, color='gray', linestyle='--', alpha=0.5)
+        ax.set_xlim(1, 32)
+    
+    plt.tight_layout()
+    plt.savefig(output_path, dpi=300, bbox_inches='tight')
+    plt.close()
+    
+    print(f"  Saved layer curves: {output_path}")
+
+
+def run_visualization(model_name: str, skip_tsne: bool = False):
+    """
+    Generate all visualizations.
+    
+    Args:
+        model_name: Model to visualize
+        skip_tsne: Skip t-SNE (requires model loading)
+    """
+    print("=" * 70)
+    print("VISUALIZATION GALLERY")
+    print("=" * 70)
+    print(f"Model: {model_name}")
+    
+    output_dir = Path("/tmp/visualizations")
+    output_dir.mkdir(parents=True, exist_ok=True)
+    
+    # Load diagnosis results
+    diagnosis_results = load_diagnosis_results(model_name, output_dir)
+    if not diagnosis_results:
+        print("ERROR: No diagnosis results found.")
+        return
+    
+    print(f"Loaded results for {len(diagnosis_results)} categories")
+    
+    model_prefix = model_name.replace('/', '_')
+    
+    # 1. Hero figure
+    print("\n1. Creating hero figure...")
+    hero_path = output_dir / f"{model_prefix}_hero_figure.png"
+    create_hero_figure(diagnosis_results, hero_path, model_name)
+    s3_upload_file(hero_path, model_name)
+    
+    # 2. Layer accuracy curves
+    print("\n2. Creating layer accuracy curves...")
+    curves_path = output_dir / f"{model_prefix}_layer_curves.png"
+    create_layer_accuracy_curves(diagnosis_results, curves_path, model_name)
+    s3_upload_file(curves_path, model_name)
+    
+    # 3. t-SNE gallery (requires model)
+    if not skip_tsne:
+        print("\n3. Creating t-SNE gallery...")
+        
+        from wisent.core.models.wisent_model import WisentModel
+        
+        print(f"  Loading model: {model_name}")
+        model = WisentModel(model_name, device="cuda")
+        
+        selected = select_representative_benchmarks(diagnosis_results, n_per_type=2)
+        print(f"  Selected benchmarks: {selected}")
+        
+        tsne_path = output_dir / f"{model_prefix}_tsne_gallery.png"
+        create_tsne_gallery(model, selected, tsne_path, model_name)
+        s3_upload_file(tsne_path, model_name)
+        
+        del model
+    else:
+        print("\n3. Skipping t-SNE gallery (--skip-tsne)")
+    
+    print("\n" + "=" * 70)
+    print("VISUALIZATION COMPLETE")
+    print("=" * 70)
+    print(f"Figures saved to: {output_dir}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Visualization gallery for RepScan")
+    parser.add_argument("--model", type=str, default="Qwen/Qwen3-8B", help="Model to visualize")
+    parser.add_argument("--skip-tsne", action="store_true", help="Skip t-SNE (doesn't require model)")
+    args = parser.parse_args()
+    
+    run_visualization(args.model, skip_tsne=args.skip_tsne)