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

wisent/core/utils/device.py

@@ -95,23 +95,19 @@ def preferred_dtype(kind: DeviceKind | None = None) -> torch.dtype:
     """
     Return the preferred dtype for model loading.
     
-    Default is float32 for consistency across all devices. This ensures steering
-    vectors trained on one device work identically on another.
+    Default is device-optimized dtype (bfloat16 on CUDA, float16 on MPS, float32 on CPU).
     
     Priority:
     1. Global override set via set_default_dtype()
     2. WISENT_DTYPE environment variable ("float32", "float16", "bfloat16", "auto")
-    3. Default: float32 (consistent across all devices)
+    3. Default: device-optimized (bfloat16 on CUDA, float16 on MPS, float32 on CPU)
     
-    To use device-optimized dtypes for better performance (at cost of cross-device
-    consistency), set WISENT_DTYPE=auto or call set_default_dtype("auto").
-       
     Example:
-        >>> preferred_dtype()  # Always float32 by default
-        torch.float32
-        >>> set_default_dtype("auto")  # Use device-optimized dtypes
-        >>> preferred_dtype()  # Now bfloat16 on CUDA, float16 on MPS
+        >>> preferred_dtype()  # bfloat16 on CUDA, float16 on MPS, float32 on CPU
         torch.bfloat16
+        >>> set_default_dtype("float32")  # Force float32 everywhere
+        >>> preferred_dtype()
+        torch.float32
     """
     # Check global override first
     if _global_dtype_override is not None:
@@ -126,8 +122,8 @@ def preferred_dtype(kind: DeviceKind | None = None) -> torch.dtype:
             return device_optimized_dtype(kind)
         return env_dtype
     
-    # Default: float32 for consistency across all devices
-    return torch.float32
+    # Default: use device-optimized dtype for best performance
+    return device_optimized_dtype(kind)
 
 
 def device_optimized_dtype(kind: DeviceKind | None = None) -> torch.dtype:
@@ -159,8 +155,14 @@ def device_optimized_dtype(kind: DeviceKind | None = None) -> torch.dtype:
 # Steering Vector dtype utilities
 # ============================================================================
 
-# Default dtype for storing steering vectors (float32 for cross-device compatibility)
-STEERING_VECTOR_DTYPE = torch.float32
+
+def steering_vector_dtype() -> torch.dtype:
+    """Return the dtype for steering vectors (uses preferred_dtype())."""
+    return preferred_dtype()
+
+
+# Legacy constant for backward compatibility - use steering_vector_dtype() instead
+STEERING_VECTOR_DTYPE = torch.float32  # Deprecated: kept for backward compat only
 
 
 def save_steering_vector(
@@ -172,10 +174,7 @@ def save_steering_vector(
     metadata: dict | None = None,
 ) -> None:
     """
-    Save a steering vector with dtype metadata for cross-device compatibility.
-    
-    Vectors are always stored in float32 for consistency across devices,
-    but the original dtype is preserved in metadata for reference.
+    Save a steering vector with dtype metadata.
     
     Args:
         path: File path to save to (.pt)
@@ -187,21 +186,22 @@ def save_steering_vector(
     """
     # Store original dtype before conversion
     original_dtype = vector.dtype
+    storage_dtype = steering_vector_dtype()
     
-    # Convert to float32 for cross-device compatibility
-    vector_f32 = vector.to(dtype=STEERING_VECTOR_DTYPE, device="cpu")
+    # Store in preferred dtype
+    vector_stored = vector.to(dtype=storage_dtype, device="cpu")
     
     save_data = {
         # Primary data
-        "steering_vector": vector_f32,
+        "steering_vector": vector_stored,
         "layer": layer,
         "model": model_name,
         "method": method,
         # Dtype metadata
         "original_dtype": str(original_dtype),
-        "storage_dtype": str(STEERING_VECTOR_DTYPE),
+        "storage_dtype": str(storage_dtype),
         # Legacy keys for backward compatibility
-        "vector": vector_f32,
+        "vector": vector_stored,
         "layer_index": layer,
     }
     
@@ -233,7 +233,8 @@ def load_steering_vector(
         - "original_dtype": Original dtype when saved
         - "metadata": Any additional metadata
     """
-    data = torch.load(path, map_location="cpu", weights_only=False)
+    target_device = device or resolve_default_device()
+    data = torch.load(path, map_location=target_device, weights_only=False)
     
     # Get the vector (support both old and new key names)
     vector = data.get("steering_vector") or data.get("vector")
@@ -241,10 +242,9 @@ def load_steering_vector(
         raise FileLoadError(file_path=str(path), reason="No steering vector found")
     
     # Determine target dtype
-    target_dtype = dtype or preferred_dtype(device)
-    target_device = device or resolve_default_device()
+    target_dtype = dtype or preferred_dtype(target_device)
     
-    # Convert to target dtype/device
+    # Convert to target dtype/device if needed
     vector = vector.to(dtype=target_dtype, device=target_device)
     
     return {

wisent/core/utils/layer_combinations.py

@@ -0,0 +1,70 @@
+"""Generate all layer combinations for hyperparameter search."""
+
+from itertools import combinations
+from math import comb
+from typing import List
+
+
+def get_layer_combinations(num_layers: int, max_combo_size: int, single_and_all_only: bool = True) -> List[List[int]]:
+    """
+    Generate layer combinations up to a maximum combination size.
+    
+    Args:
+        num_layers: Total number of layers in the model
+        max_combo_size: Maximum number of layers in a combination (e.g., 3)
+        single_and_all_only: If True, only return single layers and all layers together
+                             (skip 2-layer, 3-layer combinations). Default: True
+        
+    Returns:
+        List of layer combinations:
+        - All layers together: [0, 1, 2, ..., num_layers-1]
+        - All individual layers: [0], [1], ..., [num_layers-1]
+        - (if not single_and_all_only) All combinations of 2, 3, ..., max_combo_size layers
+    """
+    all_layers = list(range(num_layers))
+    result = []
+    
+    # All layers together (always included)
+    result.append(all_layers)
+    
+    # All individual layers
+    for layer in all_layers:
+        result.append([layer])
+    
+    # All combinations of 2, 3, ..., max_combo_size layers (unless single_and_all_only)
+    if not single_and_all_only:
+        for r in range(2, max_combo_size + 1):
+            for combo in combinations(all_layers, r):
+                result.append(list(combo))
+    
+    return result
+
+
+def get_layer_combinations_count(num_layers: int, max_combo_size: int) -> int:
+    """
+    Calculate total number of layer combinations without generating them.
+    
+    Total = 1 (all layers) + C(n,1) + C(n,2) + ... + C(n, max_combo_size)
+    """
+    total = 1  # all layers
+    for r in range(1, max_combo_size + 1):
+        total += comb(num_layers, r)
+    return total
+
+
+if __name__ == "__main__":
+    # Test with 16 layers (like Llama-3.2-1B) and max_combo_size=3
+    num_layers = 16
+    max_combo_size = 3
+    combos = get_layer_combinations(num_layers, max_combo_size)
+    
+    print(f"Model with {num_layers} layers, max_combo_size={max_combo_size}:")
+    print(f"Total combinations: {len(combos)}")
+    print(f"Expected: {get_layer_combinations_count(num_layers, max_combo_size)}")
+    print()
+    
+    print("First 20 combinations:")
+    for i, combo in enumerate(combos[:20]):
+        print(f"  {i+1}: {combo}")
+    if len(combos) > 20:
+        print(f"  ... and {len(combos) - 20} more")

wisent/examples/__init__.py

@@ -0,0 +1 @@
+# Wisent examples

wisent/examples/scripts/__init__.py

@@ -0,0 +1 @@
+# Scripts for wisent examples

wisent/examples/scripts/count_all_benchmarks.py

@@ -0,0 +1,121 @@
+#!/usr/bin/env python3
+"""Count maximum contrastive pairs for all benchmarks."""
+
+import json
+import os
+import sys
+from pathlib import Path
+
+# Set environment variables
+os.environ['HF_DATASETS_TRUST_REMOTE_CODE'] = '1'
+os.environ['HF_ALLOW_CODE_EVAL'] = '1'
+
+# Add wisent to path
+sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent))
+
+from wisent.core.data_loaders.loaders.lm_loader import LMEvalDataLoader
+from wisent.core.data_loaders.loaders.huggingface_loader import HuggingFaceDataLoader
+
+
+def load_benchmarks():
+    """Load benchmarks from central registry."""
+    from wisent.core.benchmark_registry import get_all_benchmarks
+    return get_all_benchmarks()
+
+
+def count_pairs(task_name: str) -> int:
+    """Count maximum contrastive pairs for a task."""
+    try:
+        # Determine loader type
+        hf_tasks = [
+            "math", "math_500", "aime", "hmmt", "polymath", "livemathbench",
+            "humaneval", "humaneval_plus", "mbpp", "mbpp_plus",
+            "instruct_humaneval", "apps", "conala", "concode",
+            "ds", "ds1000", "ds_1000", "mercury", "recode",
+            "multipl", "multiple_", "multipl_e",
+            "codexglue", "livecodebench",
+            "super_gpqa", "supergpqa", "hle",
+            "tag",
+            "meddialog",
+            "mmlusr"
+        ]
+
+        lm_eval_only_tasks = [
+            "minerva_math", "code_x_glue", "humaneval_infilling", "mathqa"
+        ]
+
+        if any(task_name.lower() == t or task_name.lower().startswith(t + "_") for t in lm_eval_only_tasks):
+            loader = LMEvalDataLoader()
+        elif any(task_name.lower().startswith(t) for t in hf_tasks):
+            loader = HuggingFaceDataLoader()
+        else:
+            loader = LMEvalDataLoader()
+
+        # Load with no limit to get full count
+        result = loader._load_one_task(
+            task_name=task_name,
+            split_ratio=0.8,
+            seed=42,
+            limit=None,  # No limit
+            training_limit=None,
+            testing_limit=None
+        )
+
+        train_pairs = len(result['train_qa_pairs'].pairs) if result.get('train_qa_pairs') else 0
+        test_pairs = len(result['test_qa_pairs'].pairs) if result.get('test_qa_pairs') else 0
+        total_pairs = train_pairs + test_pairs
+
+        return total_pairs
+
+    except Exception as e:
+        print(f"Error counting pairs for {task_name}: {e}", file=sys.stderr)
+        return -1
+
+
+def main():
+    benchmarks = load_benchmarks()
+
+    print(f"Counting maximum contrastive pairs for {len(benchmarks)} benchmarks...\n")
+
+    results = {}
+
+    for i, benchmark in enumerate(benchmarks, 1):
+        print(f"[{i}/{len(benchmarks)}] Counting {benchmark}...", end=" ", flush=True)
+
+        count = count_pairs(benchmark)
+        results[benchmark] = count
+
+        if count >= 0:
+            print(f"{count:,} pairs")
+        else:
+            print("ERROR")
+
+    # Save results
+    output_file = Path(__file__).parent / "benchmark_pair_counts.json"
+    with open(output_file, 'w') as f:
+        json.dump(results, f, indent=2, sort_keys=True)
+
+    print(f"\nResults saved to {output_file}")
+
+    # Print summary
+    print("\n" + "="*70)
+    print("SUMMARY")
+    print("="*70)
+
+    successful = {k: v for k, v in results.items() if v >= 0}
+    failed = {k: v for k, v in results.items() if v < 0}
+
+    print(f"Successfully counted: {len(successful)}/{len(benchmarks)}")
+    print(f"Failed: {len(failed)}/{len(benchmarks)}")
+
+    if successful:
+        total = sum(successful.values())
+        avg = total / len(successful)
+        print(f"\nTotal pairs across all benchmarks: {total:,}")
+        print(f"Average pairs per benchmark: {avg:,.0f}")
+        print(f"Max pairs: {max(successful.values()):,} ({max(successful, key=successful.get)})")
+        print(f"Min pairs: {min(successful.values()):,} ({min(successful, key=successful.get)})")
+
+
+if __name__ == "__main__":
+    main()