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

wisent/core/cli/optimize_weights.py

@@ -23,10 +23,47 @@ import os
 import re
 import tempfile
 import time
+import subprocess
 from dataclasses import dataclass
 from typing import Any, Callable
 
 import torch
+from wisent.core.utils.device import resolve_default_device
+
+
+def upload_to_s3(local_path: str, s3_bucket: str, s3_key: str) -> bool:
+    """Upload a file or directory to S3."""
+    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 as e:
+        print(f"   Warning: S3 upload failed: {e}")
+        return False
+
+
+def download_from_s3(s3_bucket: str, s3_key: str, local_path: str) -> bool:
+    """Download a file or directory from S3."""
+    try:
+        s3_path = f"s3://{s3_bucket}/{s3_key}"
+        # Check if it exists
+        check_cmd = ["aws", "s3", "ls", s3_path]
+        result = subprocess.run(check_cmd, capture_output=True)
+        if result.returncode != 0:
+            return False
+        # Download
+        if s3_key.endswith('/'):
+            cmd = ["aws", "s3", "sync", s3_path, local_path, "--quiet"]
+        else:
+            os.makedirs(os.path.dirname(local_path) or '.', exist_ok=True)
+            cmd = ["aws", "s3", "cp", s3_path, local_path, "--quiet"]
+        subprocess.run(cmd, check=True, capture_output=True)
+        return True
+    except Exception:
+        return False
 
 from wisent.core.errors import UnknownTypeError, InsufficientDataError
 
@@ -204,9 +241,19 @@ def execute_optimize_weights(args):
     # Use checkpointing if checkpoint path is provided
     checkpoint_path = getattr(args, 'checkpoint', None)
     checkpoint_interval = getattr(args, 'checkpoint_interval', 5)
+    s3_bucket = getattr(args, 's3_bucket', None)
     
-    if checkpoint_path:
-        print(f"   Checkpointing enabled: {checkpoint_path}")
+    # Generate S3 key prefix for this optimization run
+    s3_key_prefix = None
+    if s3_bucket:
+        task_name = args.task.replace(',', '_')[:50] if args.task else (args.trait or 'unknown')[:50]
+        s3_key_prefix = f"optimization-checkpoints/{task_name}/{time.strftime('%Y%m%d-%H%M%S')}"
+        print(f"   S3 bucket: {s3_bucket}")
+        print(f"   S3 key prefix: {s3_key_prefix}")
+    
+    if checkpoint_path or s3_bucket:
+        if checkpoint_path:
+            print(f"   Checkpointing enabled: {checkpoint_path}")
         print(f"   Checkpoint interval: every {checkpoint_interval} trials\n")
         result = optimizer.optimize_with_checkpointing(
             hpo_config,
@@ -214,6 +261,8 @@ def execute_optimize_weights(args):
             checkpoint_interval=checkpoint_interval,
             output_dir=args.output_dir,
             tokenizer=tokenizer,
+            s3_bucket=s3_bucket,
+            s3_key_prefix=s3_key_prefix,
         )
     else:
         result = optimizer.optimize(hpo_config)
@@ -270,6 +319,17 @@ def execute_optimize_weights(args):
     print(f"   Model saved")
     print(f"   Metadata saved to optimization_metadata.json")
 
+    # Upload to S3 if --s3-bucket is provided
+    s3_bucket = getattr(args, 's3_bucket', None)
+    if s3_bucket:
+        task_name = args.task.replace(',', '_')[:50] if args.task else (args.trait or 'unknown')[:50]
+        s3_key = f"optimization-results/{task_name}/{time.strftime('%Y%m%d-%H%M%S')}"
+        print(f"\n   Uploading results to s3://{s3_bucket}/{s3_key}/...")
+        if upload_to_s3(args.output_dir, s3_bucket, s3_key):
+            print(f"   ✓ Results uploaded to S3")
+        else:
+            print(f"   ✗ S3 upload failed")
+
     # Save all trials if requested
     if args.save_trials:
         trials_data = [
@@ -569,10 +629,10 @@ def _generate_steering_vectors(args, num_pairs: int, num_layers: int = None) ->
             # Use .pt format for train_unified_goodness output
             temp_output_pt = temp_output.replace('.json', '.pt')
 
-            # Parse layers - if 'all' or None, use None to let train_unified_goodness pick middle layer
+            # Parse layers - if 'all' or None, use None to let train_unified_goodness pick ALL layers
             layers_arg = args.layers if hasattr(args, 'layers') else None
             if layers_arg == 'all' or layers_arg is None:
-                layers_arg = None  # Will use middle layer
+                layers_arg = None  # Will use ALL layers (train_unified_goodness default)
             
             vector_args = Namespace(
                 task=args.task,  # Pass comma-separated benchmarks
@@ -602,7 +662,7 @@ def _generate_steering_vectors(args, num_pairs: int, num_layers: int = None) ->
             execute_train_unified_goodness(vector_args)
 
             # Load the .pt file
-            checkpoint = torch.load(temp_output_pt, map_location='cpu', weights_only=False)
+            checkpoint = torch.load(temp_output_pt, map_location=resolve_default_device(), weights_only=False)
 
             # Handle different checkpoint formats
             if 'all_layer_vectors' in checkpoint:
@@ -907,8 +967,8 @@ def _create_custom_evaluator(args, model_name: str) -> Callable:
             )
             response = result[0] if result else ""
             
-            # Score using custom evaluator
-            score = custom_eval(response)
+            # Score using custom evaluator - pass prompt for coherence checking
+            score = custom_eval(response, prompt=prompt_text)
             if isinstance(score, dict):
                 # Take the primary score (first value or 'score' key)
                 score = score.get('score', list(score.values())[0])

wisent/core/cli/preview_pairs.py

@@ -0,0 +1,203 @@
+"""Preview contrastive pairs from benchmarks with different extraction strategies."""
+
+import sys
+import json
+import argparse
+from typing import Optional
+
+
+def execute_preview_pairs(args):
+    """Preview contrastive pairs from a benchmark with different strategies applied."""
+    from wisent.core.contrastive_pairs.lm_eval_pairs.lm_task_pairs_generation import (
+        lm_build_contrastive_pairs,
+    )
+    from wisent.core.contrastive_pairs.huggingface_pairs.hf_extractor_manifest import HF_EXTRACTORS
+    from wisent.core.activations.extraction_strategy import (
+        ExtractionStrategy,
+        build_extraction_texts,
+        get_strategy_for_model,
+    )
+    
+    task_name = args.task_name
+    limit = args.limit or 5
+    strategies = args.strategies or ['chat_last', 'mc_balanced', 'completion_last']
+    
+    print(f"\n{'='*80}")
+    print(f"Preview Contrastive Pairs: {task_name}")
+    print(f"{'='*80}")
+    
+    # Load pairs
+    print(f"\nLoading {limit} pairs from '{task_name}'...")
+    
+    try:
+        task_name_lower = task_name.lower()
+        is_hf_task = task_name_lower in {k.lower() for k in HF_EXTRACTORS.keys()}
+        
+        if is_hf_task:
+            pairs = lm_build_contrastive_pairs(
+                task_name=task_name,
+                lm_eval_task=None,
+                limit=limit,
+            )
+        else:
+            from wisent.core.data_loaders.loaders.lm_loader import LMEvalDataLoader
+            loader = LMEvalDataLoader()
+            task_obj = loader.load_lm_eval_task(task_name)
+            
+            if isinstance(task_obj, dict):
+                if len(task_obj) != 1:
+                    keys = ", ".join(sorted(task_obj.keys()))
+                    print(f"Task '{task_name}' has subtasks: {keys}")
+                    print("Please specify a subtask.")
+                    sys.exit(1)
+                (subname, task), = task_obj.items()
+                task_name = subname
+            else:
+                task = task_obj
+            
+            pairs = lm_build_contrastive_pairs(
+                task_name=task_name,
+                lm_eval_task=task,
+                limit=limit,
+            )
+        
+        print(f"Loaded {len(pairs)} pairs\n")
+        
+    except Exception as e:
+        print(f"Error loading task: {e}")
+        sys.exit(1)
+    
+    # Mock tokenizer for preview
+    class PreviewTokenizer:
+        def apply_chat_template(self, messages, tokenize=False, add_generation_prompt=False):
+            if len(messages) == 1:
+                return f"<|user|>\n{messages[0]['content']}\n<|assistant|>\n"
+            elif len(messages) == 2:
+                return f"<|user|>\n{messages[0]['content']}\n<|assistant|>\n{messages[1]['content']}<|end|>"
+            return str(messages)
+        
+        def __call__(self, text, add_special_tokens=False):
+            return {"input_ids": text.split()}
+    
+    tokenizer = PreviewTokenizer()
+    
+    # Show pairs with strategies
+    for i, pair in enumerate(pairs):
+        print(f"\n{'='*80}")
+        print(f"PAIR {i+1}/{len(pairs)}")
+        print(f"{'='*80}")
+        
+        print(f"\n--- RAW DATA (from extractor) ---")
+        print(f"Prompt: {pair.prompt[:300]}{'...' if len(pair.prompt) > 300 else ''}")
+        print(f"Correct: {pair.positive_response.model_response[:100]}{'...' if len(pair.positive_response.model_response) > 100 else ''}")
+        print(f"Incorrect: {pair.negative_response.model_response[:100]}{'...' if len(pair.negative_response.model_response) > 100 else ''}")
+        
+        for strategy_name in strategies:
+            try:
+                strategy = ExtractionStrategy(strategy_name)
+            except ValueError:
+                print(f"\n--- {strategy_name.upper()} --- (invalid strategy)")
+                continue
+            
+            print(f"\n--- {strategy_name.upper()} ---")
+            
+            try:
+                # Build texts for positive response
+                if strategy in (ExtractionStrategy.MC_BALANCED, ExtractionStrategy.MC_COMPLETION):
+                    full_text, answer, prompt_only = build_extraction_texts(
+                        strategy,
+                        pair.prompt,
+                        pair.positive_response.model_response,
+                        tokenizer,
+                        other_response=pair.negative_response.model_response,
+                        is_positive=True,
+                        auto_convert_strategy=False,
+                    )
+                else:
+                    full_text, answer, prompt_only = build_extraction_texts(
+                        strategy,
+                        pair.prompt,
+                        pair.positive_response.model_response,
+                        tokenizer,
+                        auto_convert_strategy=False,
+                    )
+                
+                print(f"Full text (positive):")
+                print(f"  {full_text[:400]}{'...' if len(full_text) > 400 else ''}")
+                print(f"Answer token: {answer}")
+                
+            except Exception as e:
+                print(f"  Error: {e}")
+    
+    # Summary
+    print(f"\n{'='*80}")
+    print("SUMMARY")
+    print(f"{'='*80}")
+    print(f"Task: {task_name}")
+    print(f"Pairs shown: {len(pairs)}")
+    print(f"Strategies: {', '.join(strategies)}")
+    print()
+    
+    # Save to JSON if requested
+    if args.output:
+        output_data = {
+            "task_name": task_name,
+            "num_pairs": len(pairs),
+            "strategies": strategies,
+            "pairs": []
+        }
+        
+        for pair in pairs:
+            pair_data = {
+                "raw": {
+                    "prompt": pair.prompt,
+                    "correct": pair.positive_response.model_response,
+                    "incorrect": pair.negative_response.model_response,
+                },
+                "formatted": {}
+            }
+            
+            for strategy_name in strategies:
+                try:
+                    strategy = ExtractionStrategy(strategy_name)
+                    if strategy in (ExtractionStrategy.MC_BALANCED, ExtractionStrategy.MC_COMPLETION):
+                        full_text, answer, _ = build_extraction_texts(
+                            strategy, pair.prompt, pair.positive_response.model_response,
+                            tokenizer, other_response=pair.negative_response.model_response,
+                            is_positive=True, auto_convert_strategy=False,
+                        )
+                    else:
+                        full_text, answer, _ = build_extraction_texts(
+                            strategy, pair.prompt, pair.positive_response.model_response,
+                            tokenizer, auto_convert_strategy=False,
+                        )
+                    pair_data["formatted"][strategy_name] = {
+                        "full_text": full_text,
+                        "answer": answer,
+                    }
+                except Exception as e:
+                    pair_data["formatted"][strategy_name] = {"error": str(e)}
+            
+            output_data["pairs"].append(pair_data)
+        
+        with open(args.output, 'w') as f:
+            json.dump(output_data, f, indent=2)
+        print(f"Saved to: {args.output}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Preview contrastive pairs with different strategies")
+    parser.add_argument("task_name", help="Task/benchmark name (e.g., boolq, mmlu, hellaswag)")
+    parser.add_argument("--limit", "-n", type=int, default=5, help="Number of pairs to show (default: 5)")
+    parser.add_argument("--strategies", "-s", nargs="+", 
+                        default=["chat_last", "mc_balanced", "completion_last"],
+                        help="Strategies to preview")
+    parser.add_argument("--output", "-o", help="Save to JSON file")
+    parser.add_argument("--verbose", "-v", action="store_true", help="Verbose output")
+    
+    args = parser.parse_args()
+    execute_preview_pairs(args)
+
+
+if __name__ == "__main__":
+    main()

wisent/core/cli/steering_method_trainer.py

@@ -12,7 +12,8 @@ import torch
 
 from wisent.core.steering_methods.registry import SteeringMethodRegistry
 from wisent.core.activations.activations_collector import ActivationCollector
-from wisent.core.activations.core.atoms import ActivationAggregationStrategy, LayerActivations
+from wisent.core.activations.extraction_strategy import ExtractionStrategy
+from wisent.core.activations.core.atoms import LayerActivations
 from wisent.core.contrastive_pairs.core.set import ContrastivePairSet
 from wisent.core.models.core.atoms import SteeringPlan, SteeringVector
 
@@ -141,7 +142,7 @@ def collect_activations_for_pair_set(
     model,
     pair_set: ContrastivePairSet,
     layers: List[str],
-    aggregation: ActivationAggregationStrategy = ActivationAggregationStrategy.LAST_TOKEN,
+    aggregation: ExtractionStrategy = ExtractionStrategy.CHAT_LAST,
 ) -> ContrastivePairSet:
     """
     Collect activations for all pairs in a ContrastivePairSet.
@@ -155,7 +156,7 @@ def collect_activations_for_pair_set(
     Returns:
         Updated ContrastivePairSet with activations attached
     """
-    collector = ActivationCollector(model=model, store_device="cpu")
+    collector = ActivationCollector(model=model)
     
     updated_pairs = []
     for pair in pair_set.pairs:
@@ -319,7 +320,7 @@ class UnifiedSteeringTrainer:
     @property
     def collector(self) -> ActivationCollector:
         if self._collector is None:
-            self._collector = ActivationCollector(model=self.model, store_device="cpu")
+            self._collector = ActivationCollector(model=self.model)
         return self._collector
     
     def train_for_layer(
@@ -353,7 +354,7 @@ class UnifiedSteeringTrainer:
         self,
         pair_set: ContrastivePairSet,
         layers: List[str],
-        aggregation: ActivationAggregationStrategy = ActivationAggregationStrategy.LAST_TOKEN,
+        aggregation: ExtractionStrategy = ExtractionStrategy.CHAT_LAST,
     ) -> LayerActivations:
         """Collect activations and train in one step."""
         # Collect activations
@@ -551,7 +552,7 @@ def get_optimal_steering_plan(
     task_name: str,
     train_pairs: "ContrastivePairSet",
     method: str = "*",
-    aggregation: ActivationAggregationStrategy = ActivationAggregationStrategy.LAST_TOKEN,
+    aggregation: ExtractionStrategy = ExtractionStrategy.CHAT_LAST,
 ) -> Optional[Tuple["SteeringPlan", Dict[str, Any]]]:
     """
     Load optimal config and create a ready-to-use steering plan.
@@ -594,7 +595,7 @@ def get_optimal_steering_plan(
     method_name = config["method"]
     
     # Collect activations for the optimal layer
-    collector = ActivationCollector(model=model, store_device="cpu")
+    collector = ActivationCollector(model=model)
     layer_str = str(layer)
     
     pos_acts = []

wisent/core/cli/steering_search_space.py

@@ -45,11 +45,13 @@ class SensorLayerConfig(str, Enum):
 class BaseSearchSpace:
     """Base search space common to all methods."""
     
-    layers: List[int] = field(default_factory=lambda: [4, 6, 8, 10, 12])
-    strengths: List[float] = field(default_factory=lambda: [0.5, 1.0, 1.5, 2.0])
-    strategies: List[str] = field(default_factory=lambda: ["constant", "initial_only", "diminishing"])
-    token_aggregations: List[str] = field(default_factory=lambda: ["last_token", "mean_pooling"])
-    prompt_constructions: List[str] = field(default_factory=lambda: ["chat_template", "direct_completion"])
+    # layers MUST be set by get_search_space() to all layers (0 to num_layers-1)
+    # Empty default ensures it's always explicitly set
+    layers: List[int] = field(default_factory=list)
+    strengths: List[float] = field(default_factory=lambda: [0.25, 0.5, 0.75, 1.0, 1.25, 1.5, 2.0])
+    strategies: List[str] = field(default_factory=lambda: ["constant", "initial_only", "diminishing", "increasing", "gaussian"])
+    token_aggregations: List[str] = field(default_factory=lambda: ["last_token", "mean_pooling", "first_token", "max_pooling", "continuation_token"])
+    prompt_constructions: List[str] = field(default_factory=lambda: ["chat_template", "direct_completion", "multiple_choice", "role_playing", "instruction_following"])
     
     def get_total_configs(self) -> int:
         return (
@@ -320,19 +322,22 @@ def get_search_space(method_name: str, num_layers: int, quick: bool = False) ->
     """
     method = method_name.upper()
     
-    # Compute layer candidates based on model size
+    # Full search uses ALL layers
+    all_layers = list(range(num_layers))
+    
+    # Quick search uses subset of layers
     if num_layers > 20:
-        default_layers = list(range(num_layers // 2, num_layers - 2, 2))
+        quick_layers = list(range(num_layers // 2, num_layers - 2, 2))
     elif num_layers > 12:
-        default_layers = [4, 6, 8, 10, 12]
+        quick_layers = [4, 6, 8, 10, 12]
     else:
-        default_layers = list(range(2, num_layers, 2))
+        quick_layers = list(range(2, num_layers, 2))
     
     if quick:
         # Reduced search space for quick testing
         if method == "CAA":
             return CAASearchSpace(
-                layers=default_layers[:3],
+                layers=quick_layers[:3],
                 strengths=[0.5, 1.0, 1.5],
                 strategies=["constant"],
                 token_aggregations=["last_token"],
@@ -340,7 +345,7 @@ def get_search_space(method_name: str, num_layers: int, quick: bool = False) ->
             )
         elif method == "PRISM":
             return PRISMSearchSpace(
-                layers=default_layers[:3],
+                layers=quick_layers[:3],
                 strengths=[0.5, 1.0, 1.5],
                 strategies=["constant"],
                 token_aggregations=["last_token"],
@@ -380,18 +385,18 @@ def get_search_space(method_name: str, num_layers: int, quick: bool = False) ->
                 optimization_steps=[100],
             )
     
-    # Full search space
+    # Full search space - uses ALL layers
     if method == "CAA":
-        return CAASearchSpace(layers=default_layers)
+        return CAASearchSpace(layers=all_layers)
     elif method == "PRISM":
-        return PRISMSearchSpace(layers=default_layers)
+        return PRISMSearchSpace(layers=all_layers)
     elif method == "PULSE":
         return PULSESearchSpace(strengths=[0.5, 1.0, 1.5, 2.0])
     elif method == "TITAN":
         return TITANSearchSpace(strengths=[0.5, 1.0, 1.5, 2.0])
     else:
         # Default to CAA search space
-        return CAASearchSpace(layers=default_layers)
+        return CAASearchSpace(layers=all_layers)
 
 
 def get_search_space_from_args(method_name: str, args, num_layers: int) -> BaseSearchSpace: