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

wisent/core/evaluators/personalization/coherence.py

@@ -19,6 +19,47 @@ __all__ = ["evaluate_quality"]
 # Global tokenizer cache
 _tokenizer_cache = {}
 
+# Function words - the glue words of English that appear in natural text
+# Real sentences need these; gibberish often lacks them
+FUNCTION_WORDS = {
+    "the", "a", "an", "is", "are", "was", "were", "be", "been", "being",
+    "have", "has", "had", "do", "does", "did", "will", "would", "could",
+    "should", "may", "might", "must", "shall", "can", "need", "dare",
+    "ought", "used", "to", "of", "in", "for", "on", "with", "at", "by",
+    "from", "as", "into", "through", "during", "before", "after", "above",
+    "below", "between", "under", "again", "further", "then", "once",
+    "and", "but", "or", "nor", "so", "yet", "both", "either", "neither",
+    "not", "only", "own", "same", "than", "too", "very", "just", "also",
+    "now", "here", "there", "when", "where", "why", "how", "all", "each",
+    "every", "few", "more", "most", "other", "some", "such", "no",
+    "any", "i", "you", "he", "she", "it", "we", "they", "me", "him", "her",
+    "us", "them", "my", "your", "his", "its", "our", "their", "this", "that",
+    "these", "those", "what", "which", "who", "whom", "whose",
+}
+
+
+def _has_low_function_word_ratio(text: str, threshold: float = 0.15) -> bool:
+    """Check if text has suspiciously low ratio of function words.
+    
+    Natural English text typically has 30-50% function words.
+    Gibberish made of strung-together nouns/jargon has very few.
+    
+    Args:
+        text: Text to check
+        threshold: Minimum ratio of function words (default 0.15)
+        
+    Returns:
+        True if text has too few function words (likely gibberish)
+    """
+    tokens = re.findall(r'\b\w+\b', text.lower())
+    if len(tokens) < 6:
+        return False  # Too short to judge
+    
+    function_count = sum(1 for t in tokens if t in FUNCTION_WORDS)
+    ratio = function_count / len(tokens)
+    
+    return ratio < threshold
+
 
 def _get_tokenizer():
     """Get a cached tokenizer for nonsense word detection."""
@@ -137,6 +178,11 @@ def _is_gibberish(text: str) -> bool:
         if validity_ratio < 0.3:
             return True
 
+    # Check 6: Function word ratio - real English has ~30-50% function words
+    # Gibberish made of strung-together nouns/jargon has very few
+    if _has_low_function_word_ratio(text, threshold=0.15):
+        return True
+
     return False
 
 

wisent/core/hyperparameter_optimizer.py

@@ -9,8 +9,8 @@ from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_sc
 from .contrastive_pairs import ContrastivePairSet
 from .steering import SteeringMethod, SteeringType
 from .activations.activations_collector import ActivationCollector
-from .activations.core.atoms import ActivationAggregationStrategy
-from .activations.prompt_construction_strategy import PromptConstructionStrategy
+from .activations.extraction_strategy import ExtractionStrategy
+
 from wisent.core.errors import OptimizationError, NoActivationDataError, InsufficientDataError
 
 logger = logging.getLogger(__name__)
@@ -352,22 +352,22 @@ class HyperparameterOptimizer:
 
         # Map aggregation string to enum
         aggregation_map = {
-            'average': ActivationAggregationStrategy.MEAN_POOLING,
-            'first': ActivationAggregationStrategy.FIRST_TOKEN,
-            'last': ActivationAggregationStrategy.LAST_TOKEN,
-            'max': ActivationAggregationStrategy.MAX_POOLING,
+            'average': ExtractionStrategy.CHAT_MEAN,
+            'first': ExtractionStrategy.CHAT_FIRST,
+            'last': ExtractionStrategy.CHAT_LAST,
+            'max': ExtractionStrategy.CHAT_MAX_NORM,
         }
-        agg_strategy = aggregation_map.get(aggregation, ActivationAggregationStrategy.MEAN_POOLING)
+        agg_strategy = aggregation_map.get(aggregation, ExtractionStrategy.CHAT_MEAN)
 
         # Map prompt strategy string to enum
         prompt_strategy_map = {
-            'multiple_choice': PromptConstructionStrategy.MULTIPLE_CHOICE,
-            'role_playing': PromptConstructionStrategy.ROLE_PLAYING,
-            'direct_completion': PromptConstructionStrategy.DIRECT_COMPLETION,
-            'instruction_following': PromptConstructionStrategy.INSTRUCTION_FOLLOWING,
-            'chat_template': PromptConstructionStrategy.CHAT_TEMPLATE,
+            'multiple_choice': ExtractionStrategy.MC_BALANCED,
+            'role_playing': ExtractionStrategy.ROLE_PLAY,
+            'direct_completion': ExtractionStrategy.CHAT_LAST,
+            'instruction_following': ExtractionStrategy.CHAT_LAST,
+            'chat_template': ExtractionStrategy.CHAT_LAST,
         }
-        prompt_strategy = prompt_strategy_map.get(prompt_construction_strategy, PromptConstructionStrategy.CHAT_TEMPLATE)
+        prompt_strategy = prompt_strategy_map.get(prompt_construction_strategy, ExtractionStrategy.CHAT_LAST)
 
         # Create activation collector
         collector = ActivationCollector(model=model, store_device="cpu")

wisent/core/lm_eval_harness_ground_truth.py

@@ -10,7 +10,7 @@ Evaluation uses the TEST portion (20%) to ensure no data leakage with training.
 import logging
 from typing import Any, Dict
 
-from wisent.core.activations.core.atoms import ActivationAggregationStrategy
+from wisent.core.activations.extraction_strategy import ExtractionStrategy
 from wisent.core.activations.activations import Activations
 from wisent.core.layer import Layer
 from wisent.core.utils.dataset_splits import get_all_docs_from_task, create_deterministic_split
@@ -674,16 +674,12 @@ class LMEvalHarnessGroundTruth:
         }
 
     def _map_token_aggregation_to_activation_method(self, token_aggregation: str):
-        """Map token aggregation string to activation method."""
-
-        mapping = {  # TODO This should be refactor, why we use strings as Token aggregation?
-            "average": ActivationAggregationStrategy.MEAN_POOLING,
-            "mean": ActivationAggregationStrategy.MEAN_POOLING,
-            "last": ActivationAggregationStrategy.LAST_TOKEN,
-            "max": ActivationAggregationStrategy.MAX_POOLING,
-        }
-
-        return mapping.get(token_aggregation.lower(), ActivationAggregationStrategy.MEAN_POOLING)
+        """Map token aggregation string to ExtractionStrategy."""
+        # Use ExtractionStrategy directly - token_aggregation should already be a valid enum value
+        try:
+            return ExtractionStrategy(token_aggregation)
+        except ValueError:
+            return ExtractionStrategy.CHAT_LAST
 
     def _is_task_interface_task(self, task_name: str) -> bool:
         """Check if this is a TaskInterface task (not an lm-eval task)."""

wisent/core/main.py

@@ -12,6 +12,7 @@ from wisent.core.branding import print_banner
 from wisent.core.cli import execute_tasks, execute_generate_pairs_from_task, execute_generate_pairs, execute_diagnose_pairs, execute_get_activations, execute_diagnose_vectors, execute_create_steering_vector, execute_generate_vector_from_task, execute_generate_vector_from_synthetic, execute_optimize_classification, execute_optimize_steering, execute_optimize_sample_size, execute_generate_responses, execute_evaluate_responses, execute_multi_steer, execute_agent, execute_modify_weights, execute_evaluate_refusal, execute_inference_config, execute_optimization_cache, execute_optimize_weights, execute_optimize
 from wisent.core.cli.train_unified_goodness import execute_train_unified_goodness
 from wisent.core.cli.check_linearity import execute_check_linearity
+from wisent.core.cli.cluster_benchmarks import execute_cluster_benchmarks
 
 
 def _should_show_banner() -> bool:
@@ -92,6 +93,8 @@ def main():
         execute_train_unified_goodness(args)
     elif args.command == 'check-linearity':
         execute_check_linearity(args)
+    elif args.command == 'cluster-benchmarks':
+        execute_cluster_benchmarks(args)
     else:
         print(f"\n✗ Command '{args.command}' is not yet implemented")
         sys.exit(1)

wisent/core/models/wisent_model.py

@@ -154,12 +154,13 @@ class WisentModel:
         layers: list[nn.Module] = []
 
         candidates = [
-            "layers",                 
-            "model.layers",           
-            "model.decoder.layers",   
-            "transformer.h",          
+            "layers",
+            "model.layers",
+            "model.decoder.layers",
+            "transformer.h",
             "base_model.model.layers",
-            "blocks", "model.blocks", 
+            "blocks", "model.blocks",
+            "gpt_neox.layers",  # Pythia models
         ]
         for path in candidates:
             obj = m
@@ -516,7 +517,7 @@ class WisentModel:
                 return_tensors="pt",
                 padding=False,  # Single prompt, no padding needed
                 truncation=True,  # Avoid errors on long inputs
-                max_length=self.tokenizer.model_max_length  # Use model's actual limit
+                max_length=self.tokenizer.model_max_length,  # Use model's actual limit
             )
             # Move tensors to the correct device (same as _batch_encode does)
             batch = {
@@ -792,7 +793,7 @@ class WisentModel:
                 return_tensors="pt",
                 padding=False,  # Single prompt, no padding needed
                 truncation=True,  # Avoid errors on long inputs
-                max_length=self.tokenizer.model_max_length  # Use model's actual limit
+                max_length=self.tokenizer.model_max_length,  # Use model's actual limit
             )
             # Move tensors to the correct device (same as _batch_encode does)
             batch = {

wisent/core/opti/methods/opti_weights.py

@@ -240,11 +240,14 @@ class WeightsOptimizer(BaseOptimizer):
                 norm_preserve=self.config.norm_preserve,
                 verbose=False,
             )
-        elif self.config.method == "additive":
+        elif self.config.method in ("additive", "titan", "prism", "pulse"):
             # Direct additive: add steering vector directly to weight matrices
-            # This is the simplest approach that worked in manual tests
+            # This modifies weights directly, not biases, so it persists when saved
+            # Used for additive and multi-direction methods (titan/prism/pulse)
             self._apply_direct_additive(params)
         else:
+            # Default fallback - use bake_steering_with_kernel
+            # Note: This adds biases which may not load correctly for some architectures
             bake_steering_with_kernel(
                 self.model,
                 self.steering_vectors,
@@ -376,6 +379,8 @@ class WeightsOptimizer(BaseOptimizer):
         checkpoint_interval: int = 5,
         output_dir: str | None = None,
         tokenizer: Any = None,
+        s3_bucket: str | None = None,
+        s3_key_prefix: str | None = None,
     ) -> HPORun:
         """
         Run optimization with checkpointing support.
@@ -436,11 +441,20 @@ class WeightsOptimizer(BaseOptimizer):
             if checkpoint_path and trial_num % checkpoint_interval == 0:
                 self._save_checkpoint(study, checkpoint_path)
                 print(f"   [Checkpoint saved at trial {trial_num}]")
+                
+                # Upload checkpoint to S3
+                if s3_bucket and s3_key_prefix:
+                    self._upload_to_s3(checkpoint_path, s3_bucket, f"{s3_key_prefix}/checkpoint.json")
 
             # Save best model at intervals
             if output_dir and trial_num % checkpoint_interval == 0:
                 if study.best_trial is not None:
                     self._save_best_model_checkpoint(study, output_dir, tokenizer)
+                    
+                    # Upload best model checkpoint to S3
+                    if s3_bucket and s3_key_prefix:
+                        checkpoint_dir = os.path.join(output_dir, "checkpoint_best")
+                        self._upload_to_s3(checkpoint_dir, s3_bucket, f"{s3_key_prefix}/checkpoint_best/")
 
         # Run optimization with callback
         study.optimize(
@@ -521,3 +535,16 @@ class WeightsOptimizer(BaseOptimizer):
         }
         with open(os.path.join(checkpoint_dir, "checkpoint_metadata.json"), "w") as f:
             json.dump(metadata, f, indent=2)
+
+    def _upload_to_s3(self, local_path: str, s3_bucket: str, s3_key: str) -> bool:
+        """Upload a file or directory to S3."""
+        import subprocess
+        try:
+            if os.path.isdir(local_path):
+                cmd = ["aws", "s3", "sync", local_path, f"s3://{s3_bucket}/{s3_key}", "--quiet"]
+            else:
+                cmd = ["aws", "s3", "cp", local_path, f"s3://{s3_bucket}/{s3_key}", "--quiet"]
+            subprocess.run(cmd, check=True, capture_output=True)
+            return True
+        except Exception:
+            return False

wisent/core/optuna/classifier/activation_generator.py

@@ -16,7 +16,7 @@ import numpy as np
 import torch
 
 from wisent.core.activations.activations_collector import ActivationCollector
-from wisent.core.activations.core.atoms import ActivationAggregationStrategy
+from wisent.core.activations.extraction_strategy import ExtractionStrategy
 from wisent.core.activations.activations import Activations
 
 logger = logging.getLogger(__name__)
@@ -29,7 +29,7 @@ class ActivationData:
     activations: torch.Tensor
     labels: torch.Tensor
     layer: int
-    aggregation: ActivationAggregationStrategy
+    aggregation: ExtractionStrategy
     metadata: dict[str, Any]
 
     def to_numpy(self) -> tuple[np.ndarray, np.ndarray]:
@@ -102,7 +102,7 @@ class GenerationConfig:
     """Configuration for activation generation."""
 
     layer_search_range: tuple[int, int]
-    aggregation_methods: Optional[list[ActivationAggregationStrategy]] = None
+    aggregation_methods: Optional[list[ExtractionStrategy]] = None
     cache_dir: Optional[str] = None
     device: Optional[str] = None
     dtype: Optional[torch.dtype] = None  # Auto-detect if None
@@ -113,10 +113,10 @@ class GenerationConfig:
             self.cache_dir = "./activation_cache"
         if not self.aggregation_methods:
             self.aggregation_methods = [
-                ActivationAggregationStrategy.MEAN_POOLING,
-                ActivationAggregationStrategy.LAST_TOKEN,
-                ActivationAggregationStrategy.FIRST_TOKEN,
-                ActivationAggregationStrategy.MAX_POOLING,
+                ExtractionStrategy.CHAT_MEAN,
+                ExtractionStrategy.CHAT_LAST,
+                ExtractionStrategy.CHAT_FIRST,
+                ExtractionStrategy.CHAT_MAX_NORM,
             ]
 
 
@@ -239,7 +239,7 @@ class ActivationGenerator:
         return activation_data
 
     def _apply_batch_aggregation(
-        self, activations: torch.Tensor, strategy: ActivationAggregationStrategy
+        self, activations: torch.Tensor, strategy: ExtractionStrategy
     ) -> torch.Tensor:
         """
         Apply aggregation strategy to a batch of activations efficiently.
@@ -258,14 +258,16 @@ class ActivationGenerator:
             return activations
         if len(activations.shape) == 3:
             # [n_samples, n_tokens, hidden_dim] -> [n_samples, hidden_dim]
-            if strategy == ActivationAggregationStrategy.MEAN_POOLING:
+            if strategy == ExtractionStrategy.CHAT_MEAN:
                 return torch.mean(activations, dim=1)
-            if strategy == ActivationAggregationStrategy.LAST_TOKEN:
+            if strategy == ExtractionStrategy.CHAT_LAST:
                 return activations[:, -1, :]
-            if strategy == ActivationAggregationStrategy.FIRST_TOKEN:
+            if strategy == ExtractionStrategy.CHAT_FIRST:
                 return activations[:, 0, :]
-            if strategy == ActivationAggregationStrategy.MAX_POOLING:
+            if strategy == ExtractionStrategy.CHAT_MAX_NORM:
                 return torch.max(activations, dim=1)[0]
+            if strategy == ExtractionStrategy.CHAT_MEAN:
+                return torch.min(activations, dim=1)[0]
             # Default to mean pooling
             self.logger.warning(f"Unknown aggregation strategy {strategy}, using mean pooling")
             return torch.mean(activations, dim=1)

wisent/core/optuna/steering/steering_optimization.py

@@ -14,7 +14,7 @@ from typing import Any, Dict, List, Optional, Tuple
 import torch
 from tqdm import tqdm
 
-from wisent.core.activations.core import ActivationAggregationStrategy
+from wisent.core.activations import ExtractionStrategy
 from wisent.core.classifier.classifier import Classifier
 from wisent.core.contrastive_pairs.contrastive_pair import ContrastivePair
 from wisent.core.contrastive_pairs.contrastive_pair_set import ContrastivePairSet
@@ -743,24 +743,29 @@ class SteeringOptimizer:
         # Apply aggregation strategy
         if (
             aggregation_strategy == "mean_pooling"
-            or aggregation_strategy == ActivationAggregationStrategy.MEAN_POOLING.value
+            or aggregation_strategy == ExtractionStrategy.CHAT_MEAN.value
         ):
             aggregated = torch.mean(activation_tensor, dim=1)  # [1, hidden_dim]
         elif (
             aggregation_strategy == "last_token"
-            or aggregation_strategy == ActivationAggregationStrategy.LAST_TOKEN.value
+            or aggregation_strategy == ExtractionStrategy.CHAT_LAST.value
         ):
             aggregated = activation_tensor[:, -1, :]  # [1, hidden_dim]
         elif (
             aggregation_strategy == "first_token"
-            or aggregation_strategy == ActivationAggregationStrategy.FIRST_TOKEN.value
+            or aggregation_strategy == ExtractionStrategy.CHAT_FIRST.value
         ):
             aggregated = activation_tensor[:, 0, :]  # [1, hidden_dim]
         elif (
             aggregation_strategy == "max_pooling"
-            or aggregation_strategy == ActivationAggregationStrategy.MAX_POOLING.value
+            or aggregation_strategy == ExtractionStrategy.CHAT_MAX_NORM.value
         ):
             aggregated = torch.max(activation_tensor, dim=1)[0]  # [1, hidden_dim]
+        elif (
+            aggregation_strategy == "min_pooling"
+            or aggregation_strategy == ExtractionStrategy.CHAT_MEAN.value
+        ):
+            aggregated = torch.min(activation_tensor, dim=1)[0]  # [1, hidden_dim]
         else:
             # Default to mean pooling if unknown
             self.logger.warning(f"Unknown aggregation strategy {aggregation_strategy}, using mean pooling")
@@ -1029,10 +1034,10 @@ class SteeringOptimizer:
             generation_config = GenerationConfig(
                 layer_search_range=(0, 23),  # Will be auto-detected from model
                 aggregation_methods=[
-                    ActivationAggregationStrategy.MEAN_POOLING,
-                    ActivationAggregationStrategy.LAST_TOKEN,
-                    ActivationAggregationStrategy.FIRST_TOKEN,
-                    ActivationAggregationStrategy.MAX_POOLING,
+                    ExtractionStrategy.CHAT_MEAN,
+                    ExtractionStrategy.CHAT_LAST,
+                    ExtractionStrategy.CHAT_FIRST,
+                    ExtractionStrategy.CHAT_MAX_NORM,
                 ],
                 cache_dir="./cache/steering_activations",
                 device=optimization_config.device,

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_task_parser.py

@@ -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/main_parser.py

@@ -39,6 +39,7 @@ 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
 
 
 def setup_parser() -> argparse.ArgumentParser:
@@ -217,4 +218,11 @@ 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)
+
     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