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

wisent/core/geometry_search_space.py

@@ -0,0 +1,237 @@
+"""
+Configuration for geometry search space.
+
+Defines all parameters to search over when testing if a unified "goodness" 
+direction exists across benchmarks.
+
+Strategy:
+- Extract activations for ALL layers once per (benchmark, strategy) pair
+- Cache activations to disk/memory
+- Test all layer combinations from cached activations (fast, just tensor math)
+- This reduces extraction time from O(layer_combos) to O(1) per benchmark
+"""
+
+from dataclasses import dataclass, field
+from typing import List, Optional, Dict, Any
+from enum import Enum
+from pathlib import Path
+import json
+
+from wisent.core.activations.extraction_strategy import ExtractionStrategy
+from wisent.core.utils.layer_combinations import get_layer_combinations
+from wisent.core.benchmark_registry import get_all_benchmarks
+from wisent.core.activations.activation_cache import ActivationCache, CachedActivations
+
+
+@dataclass
+class GeometrySearchConfig:
+    """Configuration for a single geometry search run."""
+    
+    # Pairs settings
+    pairs_per_benchmark: int = 50
+    random_seed: int = 42
+    
+    # Layer settings
+    max_layer_combo_size: int = 3
+    
+    # Caching
+    cache_activations: bool = True
+    cache_dir: Optional[str] = None
+    
+    # Estimation
+    estimated_time_per_extraction_seconds: float = 120.0  # ~2 min per (benchmark, strategy)
+    
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "pairs_per_benchmark": self.pairs_per_benchmark,
+            "random_seed": self.random_seed,
+            "max_layer_combo_size": self.max_layer_combo_size,
+            "cache_activations": self.cache_activations,
+            "cache_dir": self.cache_dir,
+            "estimated_time_per_extraction_seconds": self.estimated_time_per_extraction_seconds,
+        }
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "GeometrySearchConfig":
+        return cls(**{k: v for k, v in data.items() if k in cls.__dataclass_fields__})
+
+
+class GeometrySearchSpace:
+    """
+    Search space configuration for geometry testing.
+    
+    Combines:
+    - Models to test
+    - Extraction strategies
+    - Layer combinations
+    - Benchmarks
+    
+    With activation caching:
+    - Extract ALL layers once per (benchmark, strategy)
+    - Test layer combinations from cache (no re-extraction needed)
+    """
+    
+    # Default models to test
+    DEFAULT_MODELS = [
+        "meta-llama/Llama-3.2-1B-Instruct",
+        "meta-llama/Llama-2-7b-chat-hf",
+        "Qwen/Qwen3-8B",
+        "openai/gpt-oss-20b",
+    ]
+    
+    # Extraction strategies for instruct models
+    INSTRUCT_STRATEGIES = [
+        ExtractionStrategy.CHAT_MEAN,
+        ExtractionStrategy.CHAT_FIRST,
+        ExtractionStrategy.CHAT_LAST,
+        ExtractionStrategy.CHAT_MAX_NORM,
+        ExtractionStrategy.CHAT_WEIGHTED,
+        ExtractionStrategy.ROLE_PLAY,
+        ExtractionStrategy.MC_BALANCED,
+    ]
+    
+    # Extraction strategies for base models
+    BASE_STRATEGIES = [
+        ExtractionStrategy.COMPLETION_LAST,
+        ExtractionStrategy.COMPLETION_MEAN,
+        ExtractionStrategy.MC_COMPLETION,
+    ]
+    
+    def __init__(
+        self,
+        models: Optional[List[str]] = None,
+        strategies: Optional[List[ExtractionStrategy]] = None,
+        benchmarks: Optional[List[str]] = None,
+        config: Optional[GeometrySearchConfig] = None,
+    ):
+        """
+        Initialize the search space.
+        
+        Args:
+            models: List of model names to test. Defaults to DEFAULT_MODELS.
+            strategies: List of extraction strategies. Defaults to INSTRUCT_STRATEGIES.
+            benchmarks: List of benchmarks. Defaults to all available benchmarks.
+            config: Search configuration (pairs, caching, etc.)
+        """
+        self.models = models or self.DEFAULT_MODELS
+        self.strategies = strategies or self.INSTRUCT_STRATEGIES
+        self.benchmarks = benchmarks or get_all_benchmarks()
+        self.config = config or GeometrySearchConfig()
+    
+    def get_layer_combinations_for_model(self, model_name: str, num_layers: int) -> List[List[int]]:
+        """
+        Get all layer combinations to test for a given model.
+        
+        Args:
+            model_name: Name of the model
+            num_layers: Number of layers in the model
+            
+        Returns:
+            List of layer combinations
+        """
+        return get_layer_combinations(num_layers, self.config.max_layer_combo_size)
+    
+    def get_extraction_count(self) -> int:
+        """
+        Calculate number of activation extractions needed (with caching).
+        
+        With caching, we extract ALL layers once per (benchmark, strategy).
+        Layer combinations are tested from cache without re-extraction.
+        
+        Returns:
+            Number of (benchmark, strategy) pairs = extraction operations
+        """
+        return len(self.benchmarks) * len(self.strategies)
+    
+    def get_total_configurations(self, num_layers: int) -> int:
+        """
+        Calculate total number of configurations to test.
+        
+        Total = strategies * layer_combos * benchmarks
+        (Layer combos are tested from cached activations)
+        """
+        from wisent.core.utils.layer_combinations import get_layer_combinations_count
+        
+        layer_combos = get_layer_combinations_count(num_layers, self.config.max_layer_combo_size)
+        return len(self.strategies) * layer_combos * len(self.benchmarks)
+    
+    def estimate_time_hours(self) -> float:
+        """
+        Estimate total time for geometry search (per model).
+        
+        With caching:
+        - Extract once per (benchmark, strategy) 
+        - Layer combo testing is fast (from cache)
+        
+        Returns:
+            Estimated hours per model
+        """
+        extractions = self.get_extraction_count()
+        seconds = extractions * self.config.estimated_time_per_extraction_seconds
+        return seconds / 3600
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Serialize to dictionary."""
+        return {
+            "models": self.models,
+            "strategies": [s.value for s in self.strategies],
+            "benchmarks": self.benchmarks,
+            "config": self.config.to_dict(),
+        }
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "GeometrySearchSpace":
+        """Deserialize from dictionary."""
+        strategies = [ExtractionStrategy(s) for s in data.get("strategies", [])]
+        config = GeometrySearchConfig.from_dict(data.get("config", {}))
+        return cls(
+            models=data.get("models"),
+            strategies=strategies if strategies else None,
+            benchmarks=data.get("benchmarks"),
+            config=config,
+        )
+    
+    def summary(self) -> str:
+        """Return a human-readable summary of the search space."""
+        lines = [
+            "Geometry Search Space:",
+            f"  Models: {len(self.models)}",
+            f"  Strategies: {len(self.strategies)}",
+            f"  Benchmarks: {len(self.benchmarks)}",
+            f"  Pairs per benchmark: {self.config.pairs_per_benchmark}",
+            f"  Max layer combo size: {self.config.max_layer_combo_size}",
+            f"  Cache activations: {self.config.cache_activations}",
+            f"",
+            f"  Extractions needed (per model): {self.get_extraction_count()}",
+            f"  Estimated time (per model): {self.estimate_time_hours():.1f} hours",
+        ]
+        return "\n".join(lines)
+    
+    def save(self, path: str) -> None:
+        """Save search space to JSON file."""
+        with open(path, "w") as f:
+            json.dump(self.to_dict(), f, indent=2)
+    
+    @classmethod
+    def load(cls, path: str) -> "GeometrySearchSpace":
+        """Load search space from JSON file."""
+        with open(path) as f:
+            return cls.from_dict(json.load(f))
+
+
+# Default search space instance
+DEFAULT_SEARCH_SPACE = GeometrySearchSpace()
+
+
+if __name__ == "__main__":
+    # Print summary of default search space
+    space = GeometrySearchSpace()
+    print(space.summary())
+    print()
+    
+    # Example with 16 layers (Llama-3.2-1B)
+    num_layers = 16
+    layer_combos = space.get_layer_combinations_for_model("test", num_layers)
+    print(f"For a {num_layers}-layer model:")
+    print(f"  Layer combinations: {len(layer_combos)}")
+    print(f"  Total configs to test: {space.get_total_configurations(num_layers)}")

wisent/core/hyperparameter_optimizer.py

@@ -370,7 +370,7 @@ class HyperparameterOptimizer:
         prompt_strategy = prompt_strategy_map.get(prompt_construction_strategy, ExtractionStrategy.CHAT_LAST)
 
         # Create activation collector
-        collector = ActivationCollector(model=model, store_device="cpu")
+        collector = ActivationCollector(model=model)
         layer_str = str(layer)
 
         # Collect activations for training pairs

wisent/core/main.py

@@ -13,6 +13,7 @@ from wisent.core.cli import execute_tasks, execute_generate_pairs_from_task, exe
 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
+from wisent.core.cli.geometry_search import execute_geometry_search
 
 
 def _should_show_banner() -> bool:
@@ -95,6 +96,8 @@ def main():
         execute_check_linearity(args)
     elif args.command == 'cluster-benchmarks':
         execute_cluster_benchmarks(args)
+    elif args.command == 'geometry-search':
+        execute_geometry_search(args)
     else:
         print(f"\n✗ Command '{args.command}' is not yet implemented")
         sys.exit(1)

wisent/core/models/core/atoms.py

@@ -7,6 +7,7 @@ import torch
 from typing import Mapping
 
 from wisent.core.errors import InvalidValueError, InvalidRangeError
+from wisent.core.utils.device import preferred_dtype
 
 if TYPE_CHECKING:
     from wisent.core.activations.core.atoms import RawActivationMap 
@@ -213,12 +214,13 @@ class SteeringPlan:
         """
         if n < 0:
             raise InvalidRangeError(param_name="n", actual=n, min_val=0)
+        dtype = preferred_dtype()
         if n == 0:
-            return torch.empty(0, dtype=torch.float32)
+            return torch.empty(0, dtype=dtype)
         if weights is None:
-            return torch.full((n,), 1.0 / n, dtype=torch.float32)
+            return torch.full((n,), 1.0 / n, dtype=dtype)
 
-        w = torch.as_tensor(weights, dtype=torch.float32)
+        w = torch.as_tensor(weights, dtype=dtype)
         if w.numel() != n:
             raise InvalidValueError(param_name="weights length", actual=w.numel(), expected=f"{n} (number of activation maps)")
         s = float(w.sum())

wisent/core/models/wisent_model.py

@@ -89,7 +89,7 @@ class WisentModel:
                 optional preloaded model (skips from_pretrained if provided).       
         """
         self.model_name = model_name
-        self.device = device or resolve_default_device()
+        self.device = resolve_default_device() if device is None or device == "auto" else device
 
         # Determine appropriate dtype and settings for the device
         load_kwargs = {

wisent/core/optuna/classifier/optuna_classifier_optimizer.py

@@ -17,7 +17,7 @@ from optuna.pruners import MedianPruner
 from optuna.samplers import TPESampler
 
 from wisent.core.classifier.classifier import Classifier
-from wisent.core.utils.device import resolve_default_device
+from wisent.core.utils.device import resolve_default_device, preferred_dtype
 from wisent.core.errors import NoActivationDataError, ClassifierCreationError
 
 from .activation_generator import ActivationData, ActivationGenerator, GenerationConfig
@@ -44,7 +44,7 @@ def get_model_dtype(model) -> torch.dtype:
         return next(model_params).dtype
     except StopIteration:
         # Fallback if no parameters found
-        return torch.float32
+        return preferred_dtype()
 
 
 logger = logging.getLogger(__name__)

wisent/core/parser_arguments/check_linearity_parser.py

@@ -1,5 +1,7 @@
 """Parser for check-linearity command."""
 
+from wisent.core.activations.extraction_strategy import ExtractionStrategy
+
 
 def setup_check_linearity_parser(parser):
     """Set up the check-linearity command parser."""
@@ -9,6 +11,14 @@ def setup_check_linearity_parser(parser):
         help='Path to JSON file containing contrastive pairs'
     )
     
+    parser.add_argument(
+        '--extraction-strategy',
+        type=str,
+        default=None,
+        choices=ExtractionStrategy.list_all(),
+        help=f'Extraction strategy to use. If not specified, tests multiple strategies. Options: {", ".join(ExtractionStrategy.list_all())}'
+    )
+    
     parser.add_argument(
         '--model',
         type=str,
@@ -19,8 +29,8 @@ def setup_check_linearity_parser(parser):
     parser.add_argument(
         '--device',
         type=str,
-        default='cuda',
-        help='Device to run model on (cuda, mps, cpu)'
+        default='auto',
+        help='Device to run model on (auto, cuda, mps, cpu)'
     )
     
     parser.add_argument(

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

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

@@ -40,6 +40,7 @@ from wisent.core.parser_arguments.train_unified_goodness_parser import setup_tra
 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:
@@ -225,4 +226,11 @@ def setup_parser() -> argparse.ArgumentParser:
     )
     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/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/methods/hyperplane.py

@@ -6,6 +6,7 @@ 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",
@@ -61,7 +62,7 @@ class HyperplaneMethod(PerLayerBaseSteeringMethod):
         clf.fit(X, y)
         
         # Use classifier weights as steering vector
-        v = torch.tensor(clf.coef_[0], dtype=torch.float32)
+        v = torch.tensor(clf.coef_[0], dtype=preferred_dtype())
         
         if bool(self.kwargs.get("normalize", True)):
             v = self._safe_l2_normalize(v)

wisent/core/synthetic/generators/nonsense_generator.py

@@ -16,16 +16,6 @@ __all__ = [
 class ProgrammaticNonsenseGenerator:
     """Generate nonsense contrastive pairs programmatically without using LLM."""
 
-    # Word list for word salad mode
-    WORD_LIST = [
-        "purple", "elephant", "calculator", "yesterday", "moon", "basket", "thinking",
-        "telephone", "mountain", "running", "quickly", "tomorrow", "happiness", "keyboard",
-        "window", "dancing", "coffee", "planet", "singing", "computer", "orange", "flying",
-        "bicycle", "dream", "ocean", "pencil", "laughing", "cloud", "table", "walking",
-        "music", "river", "chair", "jumping", "sun", "book", "swimming", "star", "door",
-        "cooking", "tree", "writing", "sky", "flower", "playing", "rain", "paper", "sleeping"
-    ]
-
     def __init__(
         self,
         nonsense_mode: str,
@@ -46,6 +36,18 @@ class ProgrammaticNonsenseGenerator:
         self.contrastive_set_name = contrastive_set_name
         self.trait_label = trait_label
         self.trait_description = trait_description
+        self._valid_words = None
+    
+    def set_tokenizer(self, tokenizer) -> None:
+        """Extract valid words from tokenizer vocabulary."""
+        vocab = tokenizer.get_vocab()
+        valid_words = []
+        for token, token_id in vocab.items():
+            decoded = tokenizer.decode([token_id])
+            clean = decoded.strip()
+            if clean.isalpha() and len(clean) > 1 and len(clean) < 15:
+                valid_words.append(clean)
+        self._valid_words = list(set(valid_words))
 
     def generate(self, num_pairs: int = 10) -> ContrastivePairSet:
         """
@@ -108,11 +110,14 @@ class ProgrammaticNonsenseGenerator:
 
     def _generate_repetitive(self) -> str:
         """Generate pathologically repetitive text."""
+        if self._valid_words is None:
+            raise ValueError("Tokenizer must be set. Call set_tokenizer() first.")
+        
         # Pick a random word or phrase
         choices = [
             random.choice(string.ascii_lowercase),  # Single letter
-            random.choice(self.WORD_LIST),  # Single word
-            ' '.join(random.sample(self.WORD_LIST, 2)),  # Two-word phrase
+            random.choice(self._valid_words),  # Single word
+            ' '.join(random.sample(self._valid_words, 2)),  # Two-word phrase
         ]
         unit = random.choice(choices)
 
@@ -121,13 +126,20 @@ class ProgrammaticNonsenseGenerator:
         return ' '.join([unit] * repetitions)
 
     def _generate_word_salad(self) -> str:
-        """Generate word salad (real words, no meaning)."""
-        num_words = random.randint(8, 15)
-        words = random.choices(self.WORD_LIST, k=num_words)
-        return ' '.join(words)
+        """Generate word salad (random tokens from tokenizer vocabulary)."""
+        num_words = random.randint(3, 10)
+        
+        if self._valid_words is not None:
+            words = random.choices(self._valid_words, k=num_words)
+            return ' '.join(words)
+        
+        raise ValueError("Tokenizer must be set to generate word salad. Call set_tokenizer() first.")
 
     def _generate_mixed(self) -> str:
         """Generate mixed nonsense (combination of all types)."""
+        if self._valid_words is None:
+            raise ValueError("Tokenizer must be set. Call set_tokenizer() first.")
+        
         components = []
 
         # Add 2-4 different types of nonsense
@@ -140,11 +152,11 @@ class ProgrammaticNonsenseGenerator:
                 length = random.randint(5, 15)
                 components.append(''.join(random.choices(string.ascii_lowercase, k=length)))
             elif mode == 'repetitive':
-                word = random.choice(self.WORD_LIST)
+                word = random.choice(self._valid_words)
                 reps = random.randint(3, 6)
                 components.append(' '.join([word] * reps))
             else:  # word_salad
                 num_words = random.randint(3, 6)
-                components.append(' '.join(random.choices(self.WORD_LIST, k=num_words)))
+                components.append(' '.join(random.choices(self._valid_words, k=num_words)))
 
         return ' '.join(components)

wisent/core/trainers/steering_trainer.py

@@ -48,8 +48,8 @@ class WisentSteeringTrainer(BaseSteeringTrainer):
         model: WisentModel to use for activation collection.
         pair_set: ContrastivePairSet with pairs to use for collection and training.
         steering_method: BaseSteeringMethod instance to use for training.
-        store_device: Device to store collected activations on (default "cpu").
-        dtype: Optional torch.dtype to cast collected activations to (default None, meaning no cast).
+        store_device: Device to store collected activations on (default: "cpu" to avoid GPU OOM).
+        dtype: Optional torch.dtype to cast collected activations to.
     """
 
     model: WisentModel