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

wisent/core/agent/diagnose/response_diagnostics.py

@@ -11,7 +11,7 @@ This module handles:
 from dataclasses import dataclass
 from typing import Any, Dict, List
 
-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.classifier.classifier import Classifier
 from wisent.core.layer import Layer
@@ -131,7 +131,7 @@ class ResponseDiagnostics:
             activations = Activations(
                 tensor=activations_tensor,
                 layer=classifier_config["layer"],
-                aggregation_strategy=ActivationAggregationStrategy.LAST_TOKEN,
+                aggregation_strategy=ExtractionStrategy.CHAT_LAST,
             )
 
             # Get features for classifier
@@ -177,7 +177,7 @@ class ResponseDiagnostics:
             activations = Activations(
                 tensor=activations_tensor,
                 layer=classifier_config["layer"],
-                aggregation_strategy=ActivationAggregationStrategy.LAST_TOKEN,
+                aggregation_strategy=ExtractionStrategy.CHAT_LAST,
             )
 
             # Get features

wisent/core/agent/diagnose.py

@@ -11,7 +11,7 @@ This module handles:
 from dataclasses import dataclass
 from typing import Any, Dict, List
 
-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.classifier.classifier import Classifier
 from wisent.core.errors import (
@@ -121,7 +121,7 @@ class ResponseDiagnostics:
             activations = Activations(
                 tensor=activations_tensor,
                 layer=classifier_config["layer"],
-                aggregation_strategy=ActivationAggregationStrategy.LAST_TOKEN,
+                aggregation_strategy=ExtractionStrategy.CHAT_LAST,
             )
 
             # Get features for classifier
@@ -161,7 +161,7 @@ class ResponseDiagnostics:
             activations = Activations(
                 tensor=activations_tensor,
                 layer=classifier_config["layer"],
-                aggregation_strategy=ActivationAggregationStrategy.LAST_TOKEN,
+                aggregation_strategy=ExtractionStrategy.CHAT_LAST,
             )
 
             # Get features

wisent/core/autonomous_agent.py

@@ -12,7 +12,7 @@ A model that can autonomously use wisent capabilities on itself:
 import asyncio
 from typing import Any, Dict, List, Optional
 
-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.models.inference_config import get_generate_kwargs
 from wisent.core.errors import MissingParameterError
@@ -333,7 +333,7 @@ class AutonomousAgent:
         activations = Activations(
             tensor=activations_tensor,
             layer=Layer(index=classifier_params.optimal_layer, type="transformer"),
-            aggregation_strategy=ActivationAggregationStrategy.LAST_TOKEN,
+            aggregation_strategy=ExtractionStrategy.CHAT_LAST,
         )
 
         features = activations.extract_features_for_classifier()

wisent/core/cli/agent/apply_steering.py

@@ -4,31 +4,31 @@ from wisent.core.models.inference_config import get_generate_kwargs
 
 
 def _map_token_aggregation(aggregation_str: str):
-    """Map string token aggregation to ActivationAggregationStrategy enum."""
-    from wisent.core.activations.core.atoms import ActivationAggregationStrategy
+    """Map string token aggregation to ExtractionStrategy."""
+    from wisent.core.activations.extraction_strategy import ExtractionStrategy
 
     mapping = {
-        "average": ActivationAggregationStrategy.MEAN_POOLING,
-        "final": ActivationAggregationStrategy.LAST_TOKEN,
-        "first": ActivationAggregationStrategy.FIRST_TOKEN,
-        "max": ActivationAggregationStrategy.MAX_POOLING,
-        "min": ActivationAggregationStrategy.MAX_POOLING,  # Note: MIN_POOLING not in enum, using MAX_POOLING
+        "average": ExtractionStrategy.CHAT_MEAN,
+        "final": ExtractionStrategy.CHAT_LAST,
+        "first": ExtractionStrategy.CHAT_FIRST,
+        "max": ExtractionStrategy.CHAT_MAX_NORM,
+        "min": ExtractionStrategy.CHAT_MEAN,
     }
-    return mapping.get(aggregation_str, ActivationAggregationStrategy.MEAN_POOLING)
+    return mapping.get(aggregation_str, ExtractionStrategy.CHAT_MEAN)
 
 
 def _map_prompt_strategy(strategy_str: str):
-    """Map string prompt strategy to PromptConstructionStrategy enum."""
-    from wisent.core.activations.prompt_construction_strategy import PromptConstructionStrategy
+    """Map string prompt strategy to ExtractionStrategy."""
+    
 
     mapping = {
-        "chat_template": PromptConstructionStrategy.CHAT_TEMPLATE,
-        "direct_completion": PromptConstructionStrategy.DIRECT_COMPLETION,
-        "instruction_following": PromptConstructionStrategy.INSTRUCTION_FOLLOWING,
-        "multiple_choice": PromptConstructionStrategy.MULTIPLE_CHOICE,
-        "role_playing": PromptConstructionStrategy.ROLE_PLAYING,
+        "chat_template": ExtractionStrategy.CHAT_LAST,
+        "direct_completion": ExtractionStrategy.CHAT_LAST,
+        "instruction_following": ExtractionStrategy.CHAT_LAST,
+        "multiple_choice": ExtractionStrategy.MC_BALANCED,
+        "role_playing": ExtractionStrategy.ROLE_PLAY,
     }
-    return mapping.get(strategy_str, PromptConstructionStrategy.CHAT_TEMPLATE)
+    return mapping.get(strategy_str, ExtractionStrategy.CHAT_LAST)
 
 
 def apply_steering_and_evaluate(
@@ -100,19 +100,17 @@ def apply_steering_and_evaluate(
     aggregation_strategy = _map_token_aggregation(token_aggregation)
     prompt_construction_strategy = _map_prompt_strategy(prompt_strategy)
 
-    # Use middle layer for steering
+    # Use specified layer for steering
     target_layers = [layer_key]
     print(f"   Target layers: {target_layers}")
 
     enriched_pairs = []
-    for i, pair in enumerate(pair_set.pairs[:10]):  # Use first 10 pairs
+    for i, pair in enumerate(pair_set.pairs):  # Use ALL pairs
         if verbose:
-            print(f"   Processing pair {i+1}/10...")
+            print(f"   Processing pair {i+1}/{len(pair_set.pairs)}...")
 
-        updated_pair = collector.collect_for_pair(
-            pair,
-            layers=target_layers,
-            aggregation=aggregation_strategy,
+        updated_pair = collector.collect(
+            pair, strategy=aggregation_strategy,
             return_full_sequence=return_full_sequence,
             normalize_layers=normalize_layers,
             prompt_strategy=prompt_construction_strategy
@@ -174,10 +172,8 @@ def apply_steering_and_evaluate(
         trait_description="evaluation"
     )
 
-    steered_evaluated_pair = collector.collect_for_pair(
-        steered_dummy_pair,
-        layers=target_layers,
-        aggregation=aggregation_strategy,
+    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

wisent/core/cli/agent/evaluate_response.py

@@ -4,31 +4,31 @@ from wisent.core.models.inference_config import get_generate_kwargs
 
 
 def _map_token_aggregation(aggregation_str: str):
-    """Map string token aggregation to ActivationAggregationStrategy enum."""
-    from wisent.core.activations.core.atoms import ActivationAggregationStrategy
+    """Map string token aggregation to ExtractionStrategy."""
+    from wisent.core.activations.extraction_strategy import ExtractionStrategy
 
     mapping = {
-        "average": ActivationAggregationStrategy.MEAN_POOLING,
-        "final": ActivationAggregationStrategy.LAST_TOKEN,
-        "first": ActivationAggregationStrategy.FIRST_TOKEN,
-        "max": ActivationAggregationStrategy.MAX_POOLING,
-        "min": ActivationAggregationStrategy.MAX_POOLING,  # Note: MIN_POOLING not in enum, using MAX_POOLING
+        "average": ExtractionStrategy.CHAT_MEAN,
+        "final": ExtractionStrategy.CHAT_LAST,
+        "first": ExtractionStrategy.CHAT_FIRST,
+        "max": ExtractionStrategy.CHAT_MAX_NORM,
+        "min": ExtractionStrategy.CHAT_MEAN,
     }
-    return mapping.get(aggregation_str, ActivationAggregationStrategy.MEAN_POOLING)
+    return mapping.get(aggregation_str, ExtractionStrategy.CHAT_MEAN)
 
 
 def _map_prompt_strategy(strategy_str: str):
-    """Map string prompt strategy to PromptConstructionStrategy enum."""
-    from wisent.core.activations.prompt_construction_strategy import PromptConstructionStrategy
+    """Map string prompt strategy to ExtractionStrategy."""
+    
 
     mapping = {
-        "chat_template": PromptConstructionStrategy.CHAT_TEMPLATE,
-        "direct_completion": PromptConstructionStrategy.DIRECT_COMPLETION,
-        "instruction_following": PromptConstructionStrategy.INSTRUCTION_FOLLOWING,
-        "multiple_choice": PromptConstructionStrategy.MULTIPLE_CHOICE,
-        "role_playing": PromptConstructionStrategy.ROLE_PLAYING,
+        "chat_template": ExtractionStrategy.CHAT_LAST,
+        "direct_completion": ExtractionStrategy.CHAT_LAST,
+        "instruction_following": ExtractionStrategy.CHAT_LAST,
+        "multiple_choice": ExtractionStrategy.MC_BALANCED,
+        "role_playing": ExtractionStrategy.ROLE_PLAY,
     }
-    return mapping.get(strategy_str, PromptConstructionStrategy.CHAT_TEMPLATE)
+    return mapping.get(strategy_str, ExtractionStrategy.CHAT_LAST)
 
 
 def evaluate_response_with_classifier(
@@ -109,10 +109,8 @@ def evaluate_response_with_classifier(
         trait_description="evaluation"
     )
 
-    evaluated_pair = collector.collect_for_pair(
-        dummy_pair,
-        layers=[layer_key],
-        aggregation=aggregation_strategy,
+    evaluated_pair = collector.collect(
+        dummy_pair, strategy=aggregation_strategy,
         return_full_sequence=return_full_sequence,
         normalize_layers=normalize_layers,
         prompt_strategy=prompt_construction_strategy

wisent/core/cli/agent/train_classifier.py

@@ -6,31 +6,31 @@ from wisent.core.errors import UnknownTypeError
 
 
 def _map_token_aggregation(aggregation_str: str):
-    """Map string token aggregation to ActivationAggregationStrategy enum."""
-    from wisent.core.activations.core.atoms import ActivationAggregationStrategy
+    """Map string token aggregation to ExtractionStrategy."""
+    from wisent.core.activations.extraction_strategy import ExtractionStrategy
 
     mapping = {
-        "average": ActivationAggregationStrategy.MEAN_POOLING,
-        "final": ActivationAggregationStrategy.LAST_TOKEN,
-        "first": ActivationAggregationStrategy.FIRST_TOKEN,
-        "max": ActivationAggregationStrategy.MAX_POOLING,
-        "min": ActivationAggregationStrategy.MAX_POOLING,  # Note: MIN_POOLING not in enum, using MAX_POOLING
+        "average": ExtractionStrategy.CHAT_MEAN,
+        "final": ExtractionStrategy.CHAT_LAST,
+        "first": ExtractionStrategy.CHAT_FIRST,
+        "max": ExtractionStrategy.CHAT_MAX_NORM,
+        "min": ExtractionStrategy.CHAT_MEAN,
     }
-    return mapping.get(aggregation_str, ActivationAggregationStrategy.MEAN_POOLING)
+    return mapping.get(aggregation_str, ExtractionStrategy.CHAT_MEAN)
 
 
 def _map_prompt_strategy(strategy_str: str):
-    """Map string prompt strategy to PromptConstructionStrategy enum."""
-    from wisent.core.activations.prompt_construction_strategy import PromptConstructionStrategy
+    """Map string prompt strategy to ExtractionStrategy."""
+    
 
     mapping = {
-        "chat_template": PromptConstructionStrategy.CHAT_TEMPLATE,
-        "direct_completion": PromptConstructionStrategy.DIRECT_COMPLETION,
-        "instruction_following": PromptConstructionStrategy.INSTRUCTION_FOLLOWING,
-        "multiple_choice": PromptConstructionStrategy.MULTIPLE_CHOICE,
-        "role_playing": PromptConstructionStrategy.ROLE_PLAYING,
+        "chat_template": ExtractionStrategy.CHAT_LAST,
+        "direct_completion": ExtractionStrategy.CHAT_LAST,
+        "instruction_following": ExtractionStrategy.CHAT_LAST,
+        "multiple_choice": ExtractionStrategy.MC_BALANCED,
+        "role_playing": ExtractionStrategy.ROLE_PLAY,
     }
-    return mapping.get(strategy_str, PromptConstructionStrategy.CHAT_TEMPLATE)
+    return mapping.get(strategy_str, ExtractionStrategy.CHAT_LAST)
 
 
 def train_classifier_on_pairs(
@@ -106,10 +106,8 @@ def train_classifier_on_pairs(
         if verbose:
             print(f"   Processing training pair {i+1}/{len(pair_set.pairs)}...")
 
-        updated_pair = collector.collect_for_pair(
-            pair,
-            layers=target_layers,
-            aggregation=aggregation_strategy,
+        updated_pair = collector.collect(
+            pair, strategy=aggregation_strategy,
             return_full_sequence=return_full_sequence,
             normalize_layers=normalize_layers,
             prompt_strategy=prompt_construction_strategy