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

wisent/core/parser_arguments/cluster_benchmarks_parser.py

@@ -0,0 +1,31 @@
+"""Parser for cluster-benchmarks command."""
+
+import argparse
+
+
+def setup_cluster_benchmarks_parser(parser: argparse.ArgumentParser) -> None:
+    """Set up arguments for the cluster-benchmarks command."""
+    parser.add_argument(
+        "--model",
+        type=str,
+        required=True,
+        help="Model name or path (e.g., meta-llama/Llama-3.2-1B-Instruct)"
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="./cluster_output",
+        help="Output directory for results (default: ./cluster_output)"
+    )
+    parser.add_argument(
+        "--pairs-per-benchmark",
+        type=int,
+        default=50,
+        help="Number of contrastive pairs per benchmark (default: 50)"
+    )
+    parser.add_argument(
+        "--device",
+        type=str,
+        default=None,
+        help="Device to use (cuda/mps/cpu). Auto-detected if not specified."
+    )

wisent/core/parser_arguments/generate_vector_from_synthetic_parser.py

@@ -40,8 +40,8 @@ def setup_generate_vector_from_synthetic_parser(parser: argparse.ArgumentParser)
     parser.add_argument(
         "--device",
         type=str,
-        default="cpu",
-        help="Device to use (e.g., 'cpu', 'cuda', 'cuda:0')"
+        default="auto",
+        help="Device to use (e.g., 'auto', 'cpu', 'cuda', 'cuda:0', 'mps')"
     )
     
     # Pair generation

wisent/core/parser_arguments/generate_vector_from_task_parser.py

@@ -46,8 +46,8 @@ def setup_generate_vector_from_task_parser(parser: argparse.ArgumentParser) -> N
     parser.add_argument(
         "--device",
         type=str,
-        default="cpu",
-        help="Device to use (e.g., 'cpu', 'cuda', 'cuda:0')"
+        default="auto",
+        help="Device to use (e.g., 'auto', 'cpu', 'cuda', 'cuda:0', 'mps')"
     )
     
     # Pair generation
@@ -101,6 +101,26 @@ def setup_generate_vector_from_task_parser(parser: argparse.ArgumentParser) -> N
         help="Do not L2-normalize steering vectors"
     )
     
+    # Universal Subspace options (PRISM/TITAN)
+    parser.add_argument(
+        "--auto-num-directions",
+        action="store_true",
+        default=False,
+        help="Automatically determine num_directions based on explained variance (PRISM/TITAN)"
+    )
+    parser.add_argument(
+        "--use-universal-basis-init",
+        action="store_true",
+        default=False,
+        help="Initialize directions from universal basis (PRISM/TITAN)"
+    )
+    parser.add_argument(
+        "--num-directions",
+        type=int,
+        default=3,
+        help="Number of steering directions for PRISM/TITAN (default: 3)"
+    )
+    
     # Intermediate file handling
     parser.add_argument(
         "--keep-intermediate",

wisent/core/parser_arguments/geometry_search_parser.py

@@ -0,0 +1,61 @@
+"""Parser for geometry-search command."""
+
+import argparse
+
+
+def setup_geometry_search_parser(parser: argparse.ArgumentParser) -> None:
+    """Set up the geometry-search command parser."""
+    parser.add_argument(
+        "--model",
+        type=str,
+        required=True,
+        help="Model name or path (e.g., meta-llama/Llama-3.2-1B-Instruct)",
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="/home/ubuntu/output/geometry_results.json",
+        help="Output path for results JSON",
+    )
+    parser.add_argument(
+        "--pairs-per-benchmark",
+        type=int,
+        default=50,
+        help="Number of pairs to sample per benchmark (default: 50)",
+    )
+    parser.add_argument(
+        "--max-layer-combo-size",
+        type=int,
+        default=3,
+        help="Maximum layers in combination (default: 3 = individual + pairs + triplets)",
+    )
+    parser.add_argument(
+        "--strategies",
+        type=str,
+        default=None,
+        help="Comma-separated list of strategies (default: all 7)",
+    )
+    parser.add_argument(
+        "--benchmarks",
+        type=str,
+        default=None,
+        help="Comma-separated list of benchmarks, or path to .txt file (default: all)",
+    )
+    parser.add_argument(
+        "--cache-dir",
+        type=str,
+        default=None,
+        help="Directory for activation cache (default: /tmp/wisent_geometry_cache_<model>)",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=42,
+        help="Random seed for reproducibility (default: 42)",
+    )
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="auto",
+        help="Device for model (auto/cuda/mps/cpu, default: auto)",
+    )

wisent/core/parser_arguments/main_parser.py

@@ -39,6 +39,8 @@ from wisent.core.parser_arguments.optimize_weights_parser import setup_optimize_
 from wisent.core.parser_arguments.train_unified_goodness_parser import setup_train_unified_goodness_parser
 from wisent.core.parser_arguments.optimize_parser import setup_optimize_parser
 from wisent.core.parser_arguments.check_linearity_parser import setup_check_linearity_parser
+from wisent.core.parser_arguments.cluster_benchmarks_parser import setup_cluster_benchmarks_parser
+from wisent.core.parser_arguments.geometry_search_parser import setup_geometry_search_parser
 
 
 def setup_parser() -> argparse.ArgumentParser:
@@ -217,4 +219,18 @@ def setup_parser() -> argparse.ArgumentParser:
     )
     setup_check_linearity_parser(check_linearity_parser)
 
+    # Cluster benchmarks command - cluster benchmarks by direction similarity
+    cluster_benchmarks_parser = subparsers.add_parser(
+        "cluster-benchmarks",
+        help="Cluster benchmarks by direction similarity with geometry analysis"
+    )
+    setup_cluster_benchmarks_parser(cluster_benchmarks_parser)
+
+    # Geometry search command - search for unified goodness direction across all benchmarks
+    geometry_search_parser = subparsers.add_parser(
+        "geometry-search",
+        help="Search for unified goodness direction across benchmarks (analyzes structure: linear/cone/orthogonal)"
+    )
+    setup_geometry_search_parser(geometry_search_parser)
+
     return parser

wisent/core/parser_arguments/optimize_steering_parser.py

@@ -68,6 +68,12 @@ def setup_steering_optimizer_parser(parser):
         default="./baseline_comparison",
         help="Directory to save baseline comparison results (default: ./baseline_comparison)",
     )
+    comprehensive_parser.add_argument(
+        "--output-dir",
+        type=str,
+        default="./optimization_results",
+        help="Directory to save optimization results (default: ./optimization_results)",
+    )
     comprehensive_parser.add_argument(
         "--show-comparisons",
         type=int,
@@ -170,32 +176,39 @@ def setup_steering_optimizer_parser(parser):
     
     # Base search space overrides
     comprehensive_parser.add_argument(
-        "--search-layers",
+        "--search-layers", "--layers",
         type=str,
         default=None,
+        dest="search_layers",
         help="Comma-separated layer indices to search (e.g., '8,10,12,14')"
     )
     comprehensive_parser.add_argument(
-        "--search-strengths",
+        "--search-strengths", "--strengths",
         type=str,
         default=None,
+        dest="search_strengths",
         help="Comma-separated strength values to search (e.g., '0.5,1.0,1.5,2.0')"
     )
     comprehensive_parser.add_argument(
-        "--search-strategies",
+        "--search-strategies", "--strategies",
         type=str,
-        nargs="+",
         default=None,
-        choices=["constant", "initial_only", "diminishing", "all_equal"],
-        help="Steering strategies to search"
+        dest="search_strategies",
+        help="Comma-separated steering strategies to search (e.g., 'constant,initial_only,diminishing,increasing,gaussian')"
     )
     comprehensive_parser.add_argument(
-        "--search-token-aggregations",
+        "--search-token-aggregations", "--token-aggregations",
         type=str,
-        nargs="+",
         default=None,
-        choices=["last_token", "mean_pooling", "first_token", "max_pooling"],
-        help="Token aggregation strategies to search"
+        dest="search_token_aggregations",
+        help="Comma-separated token aggregation strategies (e.g., 'last_token,mean_pooling,first_token,max_pooling,continuation_token,choice_token')"
+    )
+    comprehensive_parser.add_argument(
+        "--search-prompt-constructions", "--prompt-constructions",
+        type=str,
+        default=None,
+        dest="search_prompt_constructions",
+        help="Comma-separated prompt construction strategies (e.g., 'chat_template,direct_completion,instruction_following,multiple_choice,role_playing')"
     )
     
     # PRISM-specific search space
@@ -615,3 +628,97 @@ def setup_steering_optimizer_parser(parser):
         action="store_true",
         help="Save optimal parameters as default for this model/task combination",
     )
+
+    # ==========================================================================
+    # UNIVERSAL METHOD OPTIMIZER (NEW)
+    # ==========================================================================
+    # This optimizer uses the universal train(pair_set) interface that ALL
+    # steering methods implement, ensuring it works with any method including
+    # future ones.
+    
+    universal_parser = steering_subparsers.add_parser(
+        "universal",
+        help="Universal optimizer that works with ANY steering method (recommended)"
+    )
+    universal_parser.add_argument("model", type=str, help="Model name or path")
+    universal_parser.add_argument(
+        "--task",
+        type=str,
+        required=True,
+        help="Task/benchmark to optimize for (e.g., truthfulqa_generation, arc_easy)"
+    )
+    universal_parser.add_argument(
+        "--method",
+        type=str,
+        default="CAA",
+        choices=AVAILABLE_METHODS + [m.lower() for m in AVAILABLE_METHODS],
+        help=f"Steering method to optimize. Available: {', '.join(AVAILABLE_METHODS)} (default: CAA)"
+    )
+    universal_parser.add_argument(
+        "--limit",
+        type=int,
+        default=100,
+        help="Maximum samples to use (default: 100)"
+    )
+    universal_parser.add_argument(
+        "--quick",
+        action="store_true",
+        help="Use reduced search space for faster testing"
+    )
+    universal_parser.add_argument(
+        "--max-configs",
+        type=int,
+        default=None,
+        help="Maximum number of configurations to test (default: all)"
+    )
+    universal_parser.add_argument(
+        "--output-dir",
+        type=str,
+        default="./optimization_results",
+        help="Directory to save results (default: ./optimization_results)"
+    )
+    universal_parser.add_argument(
+        "--save-best-vector",
+        action="store_true",
+        help="Save the best steering vector to output directory"
+    )
+    universal_parser.add_argument("--device", type=str, default=None, help="Device to run on")
+    universal_parser.add_argument("--verbose", action="store_true", help="Enable verbose output")
+    
+    # Search space customization
+    universal_parser.add_argument(
+        "--layers",
+        type=str,
+        default=None,
+        help="Comma-separated layer indices to search (e.g., '8,10,12,14')"
+    )
+    universal_parser.add_argument(
+        "--strengths",
+        type=str,
+        default=None,
+        help="Comma-separated strength values to search (e.g., '0.5,1.0,1.5,2.0')"
+    )
+    universal_parser.add_argument(
+        "--token-aggregations",
+        type=str,
+        nargs="+",
+        default=None,
+        choices=["last_token", "mean_pooling", "first_token", "max_pooling", "continuation_token"],
+        help="Token aggregation strategies to search"
+    )
+    universal_parser.add_argument(
+        "--prompt-strategies",
+        type=str,
+        nargs="+",
+        default=None,
+        choices=["chat_template", "direct_completion", "multiple_choice", "role_playing", "instruction_following"],
+        help="Prompt construction strategies to search"
+    )
+    
+    # Method-specific parameter overrides (JSON format)
+    universal_parser.add_argument(
+        "--method-params",
+        type=str,
+        default=None,
+        help='JSON dict of method-specific parameter ranges, e.g., \'{"num_directions": [2, 3, 5]}\''
+    )

wisent/core/parser_arguments/optimize_weights_parser.py

@@ -150,6 +150,12 @@ def setup_optimize_weights_parser(parser: argparse.ArgumentParser) -> None:
         default=5,
         help="Save checkpoint and best model every N trials. Default: 5"
     )
+    parser.add_argument(
+        "--s3-bucket",
+        type=str,
+        default=None,
+        help="S3 bucket to upload results to (e.g., 'wisent-optimization-results'). Results will be uploaded on completion."
+    )
 
     # ==========================================================================
     # EVALUATION CONFIGURATION

wisent/core/parser_arguments/tasks_parser.py

@@ -155,12 +155,14 @@ def setup_tasks_parser(parser):
     parser.add_argument("--device", type=str, default=None, help="Device to run on")
     parser.add_argument("--seed", type=int, default=42, help="Random seed for reproducibility")
     parser.add_argument("--verbose", action="store_true", help="Enable verbose logging")
+    # Extraction strategy - unified approach combining prompt format and token selection
+    from wisent.core.activations.extraction_strategy import ExtractionStrategy
     parser.add_argument(
-        "--token-aggregation",
+        "--extraction-strategy",
         type=str,
-        choices=["average", "final", "first", "max", "min", "max_score"],
-        default="average",
-        help="How to aggregate token scores for classification. 'max_score' uses the highest individual token hallucination score.",
+        choices=ExtractionStrategy.list_all(),
+        default=ExtractionStrategy.default().value,
+        help=f"Extraction strategy for activations. Options: {', '.join(ExtractionStrategy.list_all())}. Default: {ExtractionStrategy.default().value}",
     )
     parser.add_argument(
         "--ground-truth-method",
@@ -385,21 +387,7 @@ def setup_tasks_parser(parser):
         help="Directory for saving/loading classifiers and vectors (default: ./models)",
     )
 
-    # Prompt construction and token targeting strategy arguments
-    parser.add_argument(
-        "--prompt-construction-strategy",
-        type=str,
-        choices=["multiple_choice", "role_playing", "direct_completion", "instruction_following", "chat_template"],
-        default="chat_template",
-        help="Strategy for constructing prompts from question-answer pairs (default: chat_template)",
-    )
-    parser.add_argument(
-        "--token-targeting-strategy",
-        type=str,
-        choices=["choice_token", "continuation_token", "last_token", "first_token", "mean_pooling", "max_pooling"],
-        default="choice_token",
-        help="Strategy for targeting tokens during activation extraction (default: choice_token)",
-    )
+
 
     # Normalization options
     parser.add_argument("--normalize-mode", action="store_true", help="Enable normalization mode (legacy flag)")

wisent/core/parser_arguments/train_unified_goodness_parser.py

@@ -32,8 +32,8 @@ def setup_train_unified_goodness_parser(parser: argparse.ArgumentParser) -> None
     parser.add_argument(
         "--device",
         type=str,
-        default="cuda",
-        help="Device to use (e.g., 'cpu', 'cuda', 'cuda:0')"
+        default="auto",
+        help="Device to use (e.g., 'auto', 'cpu', 'cuda', 'cuda:0', 'mps')"
     )
 
     # Benchmark selection

wisent/core/steering.py

@@ -477,11 +477,13 @@ class SteeringMethod:
             # Get prediction from steering method
             prediction = self.predict_proba(activation)
 
-            # Convert to tensor for loss computation
+            # Convert to tensor for loss computation (use activation's dtype)
             if not isinstance(prediction, torch.Tensor):
-                prediction = torch.tensor(prediction, dtype=torch.float32, device=self.device)
+                from wisent.core.utils.device import preferred_dtype
+                pred_dtype = activation.dtype if isinstance(activation, torch.Tensor) else preferred_dtype()
+                prediction = torch.tensor(prediction, dtype=pred_dtype, device=self.device)
 
-            target = torch.tensor(label, dtype=torch.float32, device=self.device)
+            target = torch.tensor(label, dtype=prediction.dtype, device=self.device)
 
             # Binary cross-entropy loss
             loss = F.binary_cross_entropy_with_logits(prediction.unsqueeze(0), target.unsqueeze(0))

wisent/core/steering_methods/core/atoms.py

@@ -150,5 +150,4 @@ class PerLayerBaseSteeringMethod(BaseSteeringMethod):
             raw[layer] = self.train_for_layer(pos_list, neg_list)
 
         dtype = self.kwargs.get("dtype", None)
-        agg = self.kwargs.get("activation_aggregation_strategy", None)
-        return LayerActivations(raw, activation_aggregation_strategy=agg, dtype=dtype)
+        return LayerActivations(raw, dtype=dtype)

wisent/core/steering_methods/methods/caa.py

@@ -14,7 +14,7 @@ class CAAMethod(PerLayerBaseSteeringMethod):
     """
     Contrastive Activation Additions (CAA).
     For each layer: v = mean(positives) - mean(negatives),
-    optionally L2-normalized (kwargs: normalize=True, dtype=..., activation_aggregation_strategy=...).
+    optionally L2-normalized (kwargs: normalize=True, dtype=...).
     """
     name = "caa"
     description = "Per-layer mean(pos)-mean(neg) over ContrastivePairSet."

wisent/core/steering_methods/methods/hyperplane.py

@@ -0,0 +1,75 @@
+from __future__ import annotations
+
+from typing import List
+import torch
+import numpy as np
+
+from wisent.core.steering_methods.core.atoms import PerLayerBaseSteeringMethod
+from wisent.core.errors import InsufficientDataError
+from wisent.core.utils.device import preferred_dtype
+
+__all__ = [
+    "HyperplaneMethod",
+]
+
+
+class HyperplaneMethod(PerLayerBaseSteeringMethod):
+    """
+    Hyperplane-based steering using classifier decision boundary.
+    
+    Instead of computing mean(pos) - mean(neg) like CAA, this method trains
+    a logistic regression classifier to separate positive from negative activations,
+    then uses the classifier's weight vector (hyperplane normal) as the steering vector.
+    
+    This works better when the geometry is orthogonal (each contrastive pair has
+    a unique direction) rather than linear (all pairs share a common direction).
+    In orthogonal geometry, CAA's mean difference cancels out to near-zero,
+    while the classifier can still find a separating hyperplane.
+    """
+    name = "hyperplane"
+    description = "Classifier-based steering using logistic regression decision boundary as steering vector."
+
+    def train_for_layer(self, pos_list: List[torch.Tensor], neg_list: List[torch.Tensor]) -> torch.Tensor:
+        """
+        Train hyperplane steering vector for a single layer using logistic regression.
+
+        arguments:
+            pos_list: List of positive activations (torch.Tensor) for this layer.
+            neg_list: List of negative activations (torch.Tensor) for this layer.
+            
+        returns:
+            torch.Tensor steering vector for the layer (classifier weights / hyperplane normal).
+        """
+        if not pos_list or not neg_list:
+            raise InsufficientDataError(reason="Both positive and negative lists must be non-empty.")
+        
+        pos = torch.stack([t.detach().to("cpu").float().reshape(-1) for t in pos_list], dim=0)
+        neg = torch.stack([t.detach().to("cpu").float().reshape(-1) for t in neg_list], dim=0)
+        
+        pos_np = pos.numpy()
+        neg_np = neg.numpy()
+        
+        X = np.vstack([pos_np, neg_np])
+        y = np.array([1] * len(pos_np) + [0] * len(neg_np))
+        
+        # Train logistic regression classifier
+        from sklearn.linear_model import LogisticRegression
+        
+        max_iter = int(self.kwargs.get("max_iter", 1000))
+        C = float(self.kwargs.get("C", 1.0))
+        
+        clf = LogisticRegression(max_iter=max_iter, C=C, solver="lbfgs")
+        clf.fit(X, y)
+        
+        # Use classifier weights as steering vector
+        v = torch.tensor(clf.coef_[0], dtype=preferred_dtype())
+        
+        if bool(self.kwargs.get("normalize", True)):
+            v = self._safe_l2_normalize(v)
+        
+        return v
+
+    def _safe_l2_normalize(self, v: torch.Tensor, eps: float = 1e-12) -> torch.Tensor:
+        if v.ndim != 1:
+            v = v.reshape(-1)
+        return v / (torch.linalg.norm(v) + eps)

wisent/core/steering_methods/methods/prism.py

@@ -219,8 +219,7 @@ class PRISMMethod(BaseSteeringMethod):
         primary_map: RawActivationMap = multi_result.to_single_direction_map()
         
         dtype = self.kwargs.get("dtype", None)
-        agg = self.kwargs.get("activation_aggregation_strategy", None)
-        return LayerActivations(primary_map, activation_aggregation_strategy=agg, dtype=dtype)
+        return LayerActivations(primary_map, dtype=dtype)
     
     def train_multi(self, pair_set: ContrastivePairSet) -> MultiDirectionResult:
         """

wisent/core/steering_methods/methods/pulse.py

@@ -41,15 +41,30 @@ class PULSEConfig:
     """Configuration for PULSE steering method."""
     
     # Layer configuration
-    sensor_layer: int = 15
-    """Layer index where condition gating is computed."""
+    sensor_layer: Optional[int] = None
+    """Layer index where condition gating is computed. If None, auto-computed from num_layers."""
     
-    steering_layers: List[int] = field(default_factory=lambda: [12, 13, 14, 15, 16, 17, 18])
-    """Layer indices where steering is applied."""
+    steering_layers: Optional[List[int]] = None
+    """Layer indices where steering is applied. If None, auto-computed from num_layers."""
+    
+    num_layers: Optional[int] = None
+    """Total layers in the model. Used to auto-compute steering_layers and sensor_layer."""
     
     per_layer_scaling: bool = True
     """Whether to learn/use different scaling per layer."""
     
+    def resolve_layers(self, num_layers: int) -> None:
+        """Resolve steering_layers and sensor_layer based on model's num_layers."""
+        self.num_layers = num_layers
+        if self.sensor_layer is None:
+            # 75% through the network
+            self.sensor_layer = int(num_layers * 0.75)
+        if self.steering_layers is None:
+            # Middle to late layers (50% to 85% of network)
+            start = int(num_layers * 0.5)
+            end = int(num_layers * 0.85)
+            self.steering_layers = list(range(start, end))
+    
     # Condition gating
     condition_threshold: float = 0.5
     """Threshold for condition activation (0-1)."""
@@ -188,9 +203,12 @@ class PULSEMethod(BaseSteeringMethod):
     
     def __init__(self, **kwargs: Any) -> None:
         super().__init__(**kwargs)
+        # steering_layers and sensor_layer default to None - resolved at training time
+        # based on actual num_layers in the model
         self.config = PULSEConfig(
-            sensor_layer=kwargs.get("sensor_layer", 15),
-            steering_layers=kwargs.get("steering_layers", [12, 13, 14, 15, 16, 17, 18]),
+            sensor_layer=kwargs.get("sensor_layer", None),  # Auto-resolve from num_layers
+            steering_layers=kwargs.get("steering_layers", None),  # Auto-resolve from num_layers
+            num_layers=kwargs.get("num_layers", None),
             per_layer_scaling=kwargs.get("per_layer_scaling", True),
             condition_threshold=kwargs.get("condition_threshold", 0.5),
             gate_temperature=kwargs.get("gate_temperature", 0.1),
@@ -224,8 +242,7 @@ class PULSEMethod(BaseSteeringMethod):
         
         # Return behavior vectors as LayerActivations
         dtype = self.kwargs.get("dtype", None)
-        agg = self.kwargs.get("activation_aggregation_strategy", None)
-        return LayerActivations(result.behavior_vectors, activation_aggregation_strategy=agg, dtype=dtype)
+        return LayerActivations(result.behavior_vectors, dtype=dtype)
     
     def train_pulse(
         self,
@@ -246,6 +263,20 @@ class PULSEMethod(BaseSteeringMethod):
         if condition_pairs is None:
             condition_pairs = behavior_pairs
         
+        # Detect num_layers from available data and resolve config
+        buckets = self._collect_from_set(behavior_pairs)
+        if buckets:
+            max_layer_idx = 0
+            for layer_name in buckets.keys():
+                try:
+                    layer_idx = int(str(layer_name).split("_")[-1])
+                    max_layer_idx = max(max_layer_idx, layer_idx)
+                except (ValueError, IndexError):
+                    pass
+            detected_num_layers = max_layer_idx + 1
+            if self.config.steering_layers is None or self.config.sensor_layer is None:
+                self.config.resolve_layers(detected_num_layers)
+        
         # 1. Train behavior vectors for steering layers
         behavior_vectors = self._train_behavior_vectors(behavior_pairs)