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

wisent/core/activations/extraction_strategy.py

@@ -4,15 +4,19 @@ Unified extraction strategies for activation collection.
 These strategies combine prompt construction and token extraction into a single
 unified approach, based on empirical testing of what actually works.
 
-The strategies are:
+CHAT STRATEGIES (require chat template - for instruct models):
 - chat_mean: Chat template prompt, mean of answer tokens
 - chat_first: Chat template prompt, first answer token  
 - chat_last: Chat template prompt, last token
-- chat_gen_point: Chat template prompt, token before answer (generation decision point)
 - chat_max_norm: Chat template prompt, token with max norm in answer
 - chat_weighted: Chat template prompt, position-weighted mean (earlier tokens weighted more)
 - role_play: "Behave like person who answers Q with A" format, last token
 - mc_balanced: Multiple choice with balanced A/B assignment, last token
+
+BASE MODEL STRATEGIES (no chat template - for base models like gemma-2b, gemma-9b):
+- completion_last: Direct Q+A completion, last token
+- completion_mean: Direct Q+A completion, mean of answer tokens
+- mc_completion: Multiple choice without chat template, A/B token
 """
 
 from enum import Enum
@@ -35,10 +39,7 @@ class ExtractionStrategy(str, Enum):
     """Chat template prompt with Q+A, extract first answer token."""
     
     CHAT_LAST = "chat_last"
-    """Chat template prompt with Q+A, extract last token."""
-    
-    CHAT_GEN_POINT = "chat_gen_point"
-    """Chat template prompt with Q+A, extract token before answer starts (decision point)."""
+    """Chat template prompt with Q+A, extract EOT token (has seen full answer)."""
     
     CHAT_MAX_NORM = "chat_max_norm"
     """Chat template prompt with Q+A, extract token with max norm in answer region."""
@@ -47,22 +48,34 @@ class ExtractionStrategy(str, Enum):
     """Chat template prompt with Q+A, position-weighted mean (earlier tokens weighted more)."""
     
     ROLE_PLAY = "role_play"
-    """'Behave like person who answers Q with A' format, extract last token."""
+    """'Behave like person who answers Q with A' format, extract EOT token."""
     
     MC_BALANCED = "mc_balanced"
-    """Multiple choice format with balanced A/B assignment, extract last token."""
+    """Multiple choice format with balanced A/B assignment, extract the A/B choice token."""
+    
+    # Base model strategies (no chat template required)
+    COMPLETION_LAST = "completion_last"
+    """Direct Q+A completion without chat template, extract last token. For base models."""
+    
+    COMPLETION_MEAN = "completion_mean"
+    """Direct Q+A completion without chat template, extract mean of answer tokens. For base models."""
+    
+    MC_COMPLETION = "mc_completion"
+    """Multiple choice without chat template, extract A/B token. For base models."""
 
     @property
     def description(self) -> str:
         descriptions = {
             ExtractionStrategy.CHAT_MEAN: "Chat template with mean of answer tokens",
             ExtractionStrategy.CHAT_FIRST: "Chat template with first answer token",
-            ExtractionStrategy.CHAT_LAST: "Chat template with last token",
-            ExtractionStrategy.CHAT_GEN_POINT: "Chat template with generation decision point",
+            ExtractionStrategy.CHAT_LAST: "Chat template with EOT token",
             ExtractionStrategy.CHAT_MAX_NORM: "Chat template with max-norm answer token",
             ExtractionStrategy.CHAT_WEIGHTED: "Chat template with position-weighted mean",
-            ExtractionStrategy.ROLE_PLAY: "Role-playing format with last token",
-            ExtractionStrategy.MC_BALANCED: "Balanced multiple choice with last token",
+            ExtractionStrategy.ROLE_PLAY: "Role-playing format with EOT token",
+            ExtractionStrategy.MC_BALANCED: "Balanced multiple choice with A/B token",
+            ExtractionStrategy.COMPLETION_LAST: "Direct completion with last token (base models)",
+            ExtractionStrategy.COMPLETION_MEAN: "Direct completion with mean of answer tokens (base models)",
+            ExtractionStrategy.MC_COMPLETION: "Multiple choice completion with A/B token (base models)",
         }
         return descriptions.get(self, "Unknown strategy")
     
@@ -75,6 +88,77 @@ class ExtractionStrategy(str, Enum):
     def list_all(cls) -> list[str]:
         """List all strategy names."""
         return [s.value for s in cls]
+    
+    @classmethod
+    def for_tokenizer(cls, tokenizer, prefer_mc: bool = False) -> "ExtractionStrategy":
+        """
+        Select the appropriate strategy based on whether tokenizer supports chat template.
+        
+        Args:
+            tokenizer: The tokenizer to check
+            prefer_mc: If True, prefer multiple choice strategies (mc_balanced/mc_completion)
+        
+        Returns:
+            Appropriate strategy for the tokenizer type
+        """
+        has_chat = hasattr(tokenizer, "apply_chat_template") and callable(getattr(tokenizer, "apply_chat_template"))
+        
+        if has_chat:
+            return cls.MC_BALANCED if prefer_mc else cls.CHAT_LAST
+        else:
+            return cls.MC_COMPLETION if prefer_mc else cls.COMPLETION_LAST
+    
+    @classmethod
+    def is_base_model_strategy(cls, strategy: "ExtractionStrategy") -> bool:
+        """Check if a strategy is designed for base models (no chat template)."""
+        return strategy in (cls.COMPLETION_LAST, cls.COMPLETION_MEAN, cls.MC_COMPLETION)
+    
+    @classmethod 
+    def get_equivalent_for_model_type(cls, strategy: "ExtractionStrategy", tokenizer) -> "ExtractionStrategy":
+        """
+        Get the equivalent strategy for the given tokenizer type.
+        
+        If strategy requires chat template but tokenizer doesn't have it,
+        returns the base model equivalent. And vice versa.
+        
+        Args:
+            strategy: The requested strategy
+            tokenizer: The tokenizer to check
+            
+        Returns:
+            The appropriate strategy for the tokenizer
+        """
+        has_chat = hasattr(tokenizer, "apply_chat_template") and callable(getattr(tokenizer, "apply_chat_template"))
+        is_base_strategy = cls.is_base_model_strategy(strategy)
+        
+        if has_chat and is_base_strategy:
+            # Tokenizer has chat but strategy is for base model - upgrade to chat version
+            mapping = {
+                cls.COMPLETION_LAST: cls.CHAT_LAST,
+                cls.COMPLETION_MEAN: cls.CHAT_MEAN,
+                cls.MC_COMPLETION: cls.MC_BALANCED,
+            }
+            return mapping.get(strategy, strategy)
+        
+        elif not has_chat and not is_base_strategy:
+            # Tokenizer is base model but strategy requires chat - downgrade to base version
+            mapping = {
+                cls.CHAT_LAST: cls.COMPLETION_LAST,
+                cls.CHAT_FIRST: cls.COMPLETION_LAST,
+                cls.CHAT_MEAN: cls.COMPLETION_MEAN,
+                cls.CHAT_MAX_NORM: cls.COMPLETION_LAST,
+                cls.CHAT_WEIGHTED: cls.COMPLETION_MEAN,
+                cls.ROLE_PLAY: cls.COMPLETION_LAST,
+                cls.MC_BALANCED: cls.MC_COMPLETION,
+            }
+            return mapping.get(strategy, cls.COMPLETION_LAST)
+        
+        return strategy
+
+
+def tokenizer_has_chat_template(tokenizer) -> bool:
+    """Check if tokenizer supports chat template."""
+    return hasattr(tokenizer, "apply_chat_template") and callable(getattr(tokenizer, "apply_chat_template"))
 
 
 # Random tokens for role_play strategy (deterministic based on prompt hash)
@@ -88,6 +172,7 @@ def build_extraction_texts(
     tokenizer,
     other_response: Optional[str] = None,
     is_positive: bool = True,
+    auto_convert_strategy: bool = True,
 ) -> Tuple[str, str, Optional[str]]:
     """
     Build the full text for activation extraction based on strategy.
@@ -97,8 +182,9 @@ def build_extraction_texts(
         prompt: The user prompt/question
         response: The response to extract activations for
         tokenizer: The tokenizer (needs apply_chat_template for chat strategies)
-        other_response: For mc_balanced, the other response option
-        is_positive: For mc_balanced, whether 'response' is the positive option
+        other_response: For mc_balanced/mc_completion, the other response option
+        is_positive: For mc_balanced/mc_completion, whether 'response' is the positive option
+        auto_convert_strategy: If True, automatically convert strategy to match tokenizer type
     
     Returns:
         Tuple of (full_text, answer_text, prompt_only_text)
@@ -106,31 +192,40 @@ def build_extraction_texts(
         - answer_text: The answer portion (for strategies that need it)
         - prompt_only_text: Prompt without answer (for boundary detection)
     """
+    # Auto-convert strategy if needed
+    if auto_convert_strategy:
+        original_strategy = strategy
+        strategy = ExtractionStrategy.get_equivalent_for_model_type(strategy, tokenizer)
+        if strategy != original_strategy:
+            import warnings
+            warnings.warn(
+                f"Strategy {original_strategy.value} not compatible with tokenizer, "
+                f"using {strategy.value} instead.",
+                UserWarning
+            )
     
     if strategy in (ExtractionStrategy.CHAT_MEAN, ExtractionStrategy.CHAT_FIRST,
-                    ExtractionStrategy.CHAT_LAST, ExtractionStrategy.CHAT_GEN_POINT,
-                    ExtractionStrategy.CHAT_MAX_NORM, ExtractionStrategy.CHAT_WEIGHTED):
+                    ExtractionStrategy.CHAT_LAST, ExtractionStrategy.CHAT_MAX_NORM,
+                    ExtractionStrategy.CHAT_WEIGHTED):
         # All chat_* strategies use the same prompt construction
-        if hasattr(tokenizer, "apply_chat_template"):
-            try:
-                prompt_only = tokenizer.apply_chat_template(
-                    [{"role": "user", "content": prompt}],
-                    tokenize=False,
-                    add_generation_prompt=True,
-                )
-                full_text = tokenizer.apply_chat_template(
-                    [{"role": "user", "content": prompt},
-                     {"role": "assistant", "content": response}],
-                    tokenize=False,
-                    add_generation_prompt=False,
-                )
-            except (ValueError, KeyError):
-                # Fallback for models without chat template
-                prompt_only = prompt
-                full_text = f"{prompt} {response}"
-        else:
-            prompt_only = prompt
-            full_text = f"{prompt} {response}"
+        if not hasattr(tokenizer, "apply_chat_template"):
+            raise ValueError(
+                f"Strategy {strategy.value} requires a tokenizer with apply_chat_template. "
+                f"Base models without chat templates are not supported for chat_* strategies. "
+                f"Use completion_last, completion_mean, or mc_completion instead."
+            )
+        
+        prompt_only = tokenizer.apply_chat_template(
+            [{"role": "user", "content": prompt}],
+            tokenize=False,
+            add_generation_prompt=True,
+        )
+        full_text = tokenizer.apply_chat_template(
+            [{"role": "user", "content": prompt},
+             {"role": "assistant", "content": response}],
+            tokenize=False,
+            add_generation_prompt=False,
+        )
         
         return full_text, response, prompt_only
     
@@ -139,25 +234,23 @@ def build_extraction_texts(
         random_token = ROLE_PLAY_TOKENS[hash(prompt) % len(ROLE_PLAY_TOKENS)]
         instruction = f"Behave like a person that would answer {prompt} with {response}"
         
-        if hasattr(tokenizer, "apply_chat_template"):
-            try:
-                prompt_only = tokenizer.apply_chat_template(
-                    [{"role": "user", "content": instruction}],
-                    tokenize=False,
-                    add_generation_prompt=True,
-                )
-                full_text = tokenizer.apply_chat_template(
-                    [{"role": "user", "content": instruction},
-                     {"role": "assistant", "content": random_token}],
-                    tokenize=False,
-                    add_generation_prompt=False,
-                )
-            except (ValueError, KeyError):
-                prompt_only = instruction
-                full_text = f"{instruction} {random_token}"
-        else:
-            prompt_only = instruction
-            full_text = f"{instruction} {random_token}"
+        if not hasattr(tokenizer, "apply_chat_template"):
+            raise ValueError(
+                f"Strategy {strategy.value} requires a tokenizer with apply_chat_template. "
+                f"Use completion_last or mc_completion for base models."
+            )
+        
+        prompt_only = tokenizer.apply_chat_template(
+            [{"role": "user", "content": instruction}],
+            tokenize=False,
+            add_generation_prompt=True,
+        )
+        full_text = tokenizer.apply_chat_template(
+            [{"role": "user", "content": instruction},
+             {"role": "assistant", "content": random_token}],
+            tokenize=False,
+            add_generation_prompt=False,
+        )
         
         return full_text, random_token, prompt_only
     
@@ -188,28 +281,66 @@ def build_extraction_texts(
                 option_b = response[:200]        # negative
                 answer = "B"
         
-        mc_prompt = f"Which is correct?\nA. {option_a}\nB. {option_b}\nAnswer:"
+        mc_prompt = f"Question: {prompt}\n\nWhich is correct?\nA. {option_a}\nB. {option_b}\nAnswer:"
+        
+        if not hasattr(tokenizer, "apply_chat_template"):
+            raise ValueError(
+                f"Strategy {strategy.value} requires a tokenizer with apply_chat_template. "
+                f"Use mc_completion for base models."
+            )
+        
+        prompt_only = tokenizer.apply_chat_template(
+            [{"role": "user", "content": mc_prompt}],
+            tokenize=False,
+            add_generation_prompt=True,
+        )
+        full_text = tokenizer.apply_chat_template(
+            [{"role": "user", "content": mc_prompt},
+             {"role": "assistant", "content": answer}],
+            tokenize=False,
+            add_generation_prompt=False,
+        )
+        
+        return full_text, answer, prompt_only
+    
+    elif strategy in (ExtractionStrategy.COMPLETION_LAST, ExtractionStrategy.COMPLETION_MEAN):
+        # Base model strategies - direct Q+A without chat template
+        # Format: "Q: {prompt}\nA: {response}"
+        prompt_only = f"Q: {prompt}\nA:"
+        full_text = f"Q: {prompt}\nA: {response}"
+        return full_text, response, prompt_only
+    
+    elif strategy == ExtractionStrategy.MC_COMPLETION:
+        # Multiple choice for base models - no chat template
+        if other_response is None:
+            raise ValueError("MC_COMPLETION strategy requires other_response")
+        
+        # Deterministic "random" based on prompt - same for both pos and neg of a pair
+        pos_goes_in_b = hash(prompt) % 2 == 0
         
-        if hasattr(tokenizer, "apply_chat_template"):
-            try:
-                prompt_only = tokenizer.apply_chat_template(
-                    [{"role": "user", "content": mc_prompt}],
-                    tokenize=False,
-                    add_generation_prompt=True,
-                )
-                full_text = tokenizer.apply_chat_template(
-                    [{"role": "user", "content": mc_prompt},
-                     {"role": "assistant", "content": answer}],
-                    tokenize=False,
-                    add_generation_prompt=False,
-                )
-            except (ValueError, KeyError):
-                prompt_only = mc_prompt
-                full_text = f"{mc_prompt} {answer}"
+        if is_positive:
+            if pos_goes_in_b:
+                option_a = other_response[:200]
+                option_b = response[:200]
+                answer = "B"
+            else:
+                option_a = response[:200]
+                option_b = other_response[:200]
+                answer = "A"
         else:
-            prompt_only = mc_prompt
-            full_text = f"{mc_prompt} {answer}"
+            if pos_goes_in_b:
+                option_a = response[:200]
+                option_b = other_response[:200]
+                answer = "A"
+            else:
+                option_a = other_response[:200]
+                option_b = response[:200]
+                answer = "B"
         
+        mc_prompt = f"Question: {prompt}\n\nWhich is correct?\nA. {option_a}\nB. {option_b}\nAnswer:"
+        
+        prompt_only = mc_prompt
+        full_text = f"{mc_prompt} {answer}"
         return full_text, answer, prompt_only
     
     else:
@@ -243,6 +374,7 @@ def extract_activation(
     num_answer_tokens = len(answer_tokens)
     
     if strategy == ExtractionStrategy.CHAT_LAST:
+        # EOT token - has seen the entire answer, best performance
         return hidden_states[-1]
     
     elif strategy == ExtractionStrategy.CHAT_FIRST:
@@ -257,11 +389,6 @@ def extract_activation(
             return answer_hidden.mean(dim=0)
         return hidden_states[-1]
     
-    elif strategy == ExtractionStrategy.CHAT_GEN_POINT:
-        # Last token before answer starts (decision point)
-        gen_point_idx = max(0, seq_len - num_answer_tokens - 2)
-        return hidden_states[gen_point_idx]
-    
     elif strategy == ExtractionStrategy.CHAT_MAX_NORM:
         # Token with max norm in answer region
         if num_answer_tokens > 0 and seq_len > num_answer_tokens:
@@ -275,18 +402,36 @@ def extract_activation(
         # Position-weighted mean (earlier tokens weighted more)
         if num_answer_tokens > 0 and seq_len > num_answer_tokens:
             answer_hidden = hidden_states[-num_answer_tokens-1:-1]
-            weights = torch.exp(-torch.arange(answer_hidden.shape[0], dtype=torch.float32, device=answer_hidden.device) * 0.5)
+            weights = torch.exp(-torch.arange(answer_hidden.shape[0], dtype=answer_hidden.dtype, device=answer_hidden.device) * 0.5)
             weights = weights / weights.sum()
             return (answer_hidden * weights.unsqueeze(1)).sum(dim=0)
         return hidden_states[-1]
     
-    elif strategy in (ExtractionStrategy.ROLE_PLAY, ExtractionStrategy.MC_BALANCED):
-        # Both use last token
+    elif strategy == ExtractionStrategy.ROLE_PLAY:
+        # EOT token - slightly better than answer word (65% vs 64%)
         return hidden_states[-1]
     
-    else:
-        # Default fallback
+    elif strategy == ExtractionStrategy.MC_BALANCED:
+        # Answer token (A/B) - better than EOT (64% vs 56%)
+        return hidden_states[-2]
+    
+    elif strategy == ExtractionStrategy.COMPLETION_LAST:
+        # Last token for base model completion
         return hidden_states[-1]
+    
+    elif strategy == ExtractionStrategy.COMPLETION_MEAN:
+        # Mean of answer tokens for base model completion
+        if num_answer_tokens > 0 and seq_len > num_answer_tokens:
+            answer_hidden = hidden_states[-num_answer_tokens:]
+            return answer_hidden.mean(dim=0)
+        return hidden_states[-1]
+    
+    elif strategy == ExtractionStrategy.MC_COMPLETION:
+        # A/B token for base model MC (last token is the answer)
+        return hidden_states[-1]
+    
+    else:
+        raise ValueError(f"Unknown extraction strategy: {strategy}")
 
 
 def add_extraction_strategy_args(parser: argparse.ArgumentParser) -> None:
@@ -306,3 +451,30 @@ def add_extraction_strategy_args(parser: argparse.ArgumentParser) -> None:
         choices=ExtractionStrategy.list_all(),
         help=f"Extraction strategy for activations. Options: {', '.join(ExtractionStrategy.list_all())}. Default: {ExtractionStrategy.default().value}",
     )
+
+
+def get_strategy_for_model(tokenizer, prefer_mc: bool = False) -> ExtractionStrategy:
+    """
+    Get the best extraction strategy for a given tokenizer.
+    
+    Automatically detects if tokenizer has chat template and returns
+    the appropriate strategy.
+    
+    Args:
+        tokenizer: The tokenizer to check
+        prefer_mc: If True, prefer multiple choice strategies
+        
+    Returns:
+        ExtractionStrategy appropriate for the tokenizer
+        
+    Example:
+        >>> from transformers import AutoTokenizer
+        >>> tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b")
+        >>> strategy = get_strategy_for_model(tokenizer)
+        >>> print(strategy)  # completion_last (base model)
+        
+        >>> tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B-Instruct")
+        >>> strategy = get_strategy_for_model(tokenizer)
+        >>> print(strategy)  # chat_last (instruct model)
+    """
+    return ExtractionStrategy.for_tokenizer(tokenizer, prefer_mc=prefer_mc)

wisent/core/classifiers/classifiers/core/atoms.py

@@ -14,6 +14,7 @@ import numpy as np
 
 from torch.nn.modules.loss import _Loss
 from wisent.core.errors import DuplicateNameError, InvalidRangeError, UnknownTypeError
+from wisent.core.utils.device import preferred_dtype
 
 __all__ = [
     "ClassifierTrainConfig",
@@ -164,13 +165,13 @@ class BaseClassifier(ABC):
         self,
         threshold: float = 0.5,
         device: str | None = None,
-        dtype: torch.dtype = torch.float32,
+        dtype: torch.dtype | None = None,
     ) -> None:
         if not 0.0 <= threshold <= 1.0:
             raise InvalidRangeError(param_name="threshold", actual=threshold, min_val=0.0, max_val=1.0)
         self.threshold = threshold
         self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
-        self.dtype = torch.float32 if self.device == "mps" else dtype
+        self.dtype = dtype if dtype is not None else preferred_dtype(self.device)
         self.model = None
 
     @abstractmethod

wisent/core/cli/__init__.py

@@ -22,5 +22,6 @@ from .inference_config_cli import execute_inference_config
 from .optimization_cache import execute_optimization_cache
 from .optimize_weights import execute_optimize_weights
 from .optimize import execute_optimize
+from .geometry_search import execute_geometry_search
 
-__all__ = ['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']
+__all__ = ['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', 'execute_geometry_search']

wisent/core/cli/agent/apply_steering.py

@@ -19,7 +19,7 @@ def _map_token_aggregation(aggregation_str: str):
 
 def _map_prompt_strategy(strategy_str: str):
     """Map string prompt strategy to ExtractionStrategy."""
-    
+    from wisent.core.activations.extraction_strategy import ExtractionStrategy
 
     mapping = {
         "chat_template": ExtractionStrategy.CHAT_LAST,
@@ -111,9 +111,8 @@ def apply_steering_and_evaluate(
 
         updated_pair = collector.collect(
             pair, strategy=aggregation_strategy,
-            return_full_sequence=return_full_sequence,
-            normalize_layers=normalize_layers,
-            prompt_strategy=prompt_construction_strategy
+            layers=target_layers,
+            normalize=normalize_layers
         )
         enriched_pairs.append(updated_pair)
 
@@ -174,9 +173,8 @@ def apply_steering_and_evaluate(
 
     steered_evaluated_pair = collector.collect(
         steered_dummy_pair, strategy=aggregation_strategy,
-        return_full_sequence=return_full_sequence,
-        normalize_layers=normalize_layers,
-        prompt_strategy=prompt_construction_strategy
+        layers=target_layers,
+        normalize=normalize_layers
     )
 
     steered_quality = 0.0

wisent/core/cli/agent/train_classifier.py

@@ -1,8 +1,20 @@
 """Train classifier on contrastive pairs for agent."""
 
 import numpy as np
+import torch
 from wisent.core.classifiers.classifiers.core.atoms import ClassifierTrainReport
 from wisent.core.errors import UnknownTypeError
+from wisent.core.utils.device import preferred_dtype
+
+
+def _torch_dtype_to_numpy(torch_dtype: torch.dtype):
+    """Convert torch dtype to numpy dtype."""
+    mapping = {
+        torch.float32: np.float32,
+        torch.float16: np.float16,
+        torch.bfloat16: np.float32,  # numpy doesn't support bfloat16, use float32
+    }
+    return mapping.get(torch_dtype, np.float32)
 
 
 def _map_token_aggregation(aggregation_str: str):
@@ -21,7 +33,7 @@ def _map_token_aggregation(aggregation_str: str):
 
 def _map_prompt_strategy(strategy_str: str):
     """Map string prompt strategy to ExtractionStrategy."""
-    
+    from wisent.core.activations.extraction_strategy import ExtractionStrategy
 
     mapping = {
         "chat_template": ExtractionStrategy.CHAT_LAST,
@@ -97,7 +109,7 @@ def train_classifier_on_pairs(
     prompt_construction_strategy = _map_prompt_strategy(prompt_strategy)
 
     # Collect activations for all pairs
-    collector = ActivationCollector(model=model, store_device="cpu")
+    collector = ActivationCollector(model=model)
     target_layers = [str(target_layer)]
     layer_key = target_layers[0]
 
@@ -108,9 +120,8 @@ def train_classifier_on_pairs(
 
         updated_pair = collector.collect(
             pair, strategy=aggregation_strategy,
-            return_full_sequence=return_full_sequence,
-            normalize_layers=normalize_layers,
-            prompt_strategy=prompt_construction_strategy
+            layers=[str(target_layer)],
+            normalize=normalize_layers
         )
         enriched_training_pairs.append(updated_pair)
 
@@ -133,8 +144,9 @@ def train_classifier_on_pairs(
             X_list.append(neg_act.cpu().numpy())
             y_list.append(0.0)
 
-    X_train = np.array(X_list, dtype=np.float32)
-    y_train = np.array(y_list, dtype=np.float32)
+    np_dtype = _torch_dtype_to_numpy(preferred_dtype())
+    X_train = np.array(X_list, dtype=np_dtype)
+    y_train = np.array(y_list, dtype=np_dtype)
 
     print(f"   Training data: {X_train.shape[0]} samples, {X_train.shape[1]} features")
 

wisent/core/cli/check_linearity.py

@@ -31,6 +31,7 @@ def execute_check_linearity(args):
     from wisent.core.models.wisent_model import WisentModel
     from wisent.core.contrastive_pairs.core.pair import ContrastivePair
     from wisent.core.contrastive_pairs.core.response import PositiveResponse, NegativeResponse
+    from wisent.core.activations.extraction_strategy import ExtractionStrategy
     
     # Build ContrastivePair objects
     pairs = []
@@ -72,6 +73,10 @@ def execute_check_linearity(args):
     if args.layers:
         config.layers_to_test = [int(l) for l in args.layers.split(',')]
     
+    if args.extraction_strategy:
+        config.extraction_strategies = [ExtractionStrategy(args.extraction_strategy)]
+        print(f"Using extraction strategy: {args.extraction_strategy}")
+    
     # Run check
     print("\nRunning linearity check...")
     result = check_linearity(pairs, model, config)
@@ -110,12 +115,39 @@ def execute_check_linearity(args):
         
         sorted_results = sorted(result.all_results, key=lambda x: x['linear_score'], reverse=True)
         
-        print(f"{'Linear':<8} {'d':<8} {'Layer':<6} {'Prompt':<25} {'Aggregation':<15} {'Norm'}")
-        print("-" * 80)
+        print(f"{'Linear':<8} {'d':<8} {'Layer':<6} {'Strategy':<20} {'Structure':<12} {'Norm'}")
+        print("-" * 70)
         
         for r in sorted_results[:20]:
             print(f"{r['linear_score']:<8.3f} {r['cohens_d']:<8.2f} {r['layer']:<6} "
-                  f"{r['prompt_strategy']:<25} {r['aggregation']:<15} {r['normalize']}")
+                  f"{r['extraction_strategy']:<20} {r['best_structure']:<12} {r['normalize']}")
+        
+        # Show best result for each structure type
+        if sorted_results and 'all_structure_scores' in sorted_results[0]:
+            print(f"\n{'='*60}")
+            print("BEST SCORE PER STRUCTURE TYPE")
+            print(f"{'='*60}")
+            
+            # Collect best score for each structure across all configs
+            best_per_structure = {}
+            for r in result.all_results:
+                if 'all_structure_scores' not in r:
+                    continue
+                for struct_name, data in r['all_structure_scores'].items():
+                    score = data['score']
+                    if struct_name not in best_per_structure or score > best_per_structure[struct_name]['score']:
+                        best_per_structure[struct_name] = {
+                            'score': score,
+                            'confidence': data['confidence'],
+                            'layer': r['layer'],
+                            'strategy': r['extraction_strategy'],
+                        }
+            
+            print(f"{'Structure':<12} {'Score':<8} {'Conf':<8} {'Layer':<6} {'Strategy'}")
+            print("-" * 55)
+            sorted_structs = sorted(best_per_structure.items(), key=lambda x: x[1]['score'], reverse=True)
+            for name, data in sorted_structs:
+                print(f"{name:<12} {data['score']:<8.3f} {data['confidence']:<8.3f} {data['layer']:<6} {data['strategy']}")
     
     # Exit code based on verdict
     if result.verdict.value == "linear":