wisent 0.1.1__py3-none-any.whl → 0.5.1__py3-none-any.whl

wisent/core/mixed_benchmark_sampler.py

@@ -0,0 +1,364 @@
+"""
+Mixed Benchmark Sampler for tag-based random sampling across multiple benchmarks.
+
+This module enables training and evaluation on random samples from multiple benchmarks
+that share common tags (e.g., 'coding', 'reasoning', 'math').
+"""
+
+import random
+import logging
+from typing import List, Dict, Any, Optional, Set, Tuple
+from dataclasses import dataclass
+from collections import defaultdict
+
+# Suppress BigCode debug output
+import builtins
+_original_print = getattr(builtins, '_original_print', builtins.print)
+
+def _quiet_print(*args, **kwargs):
+    """Filter out BigCode debug messages."""
+    message = ' '.join(str(arg) for arg in args)
+    if any(x in message for x in ['DEBUG', 'Available tasks:', 'ERROR extracting', 'bigcode_eval']):
+        return
+    _original_print(*args, **kwargs)
+
+# Store original print and patch
+builtins._original_print = builtins.print
+builtins.print = _quiet_print
+
+try:
+    from .lm_harness_integration.only_benchmarks import CORE_BENCHMARKS
+except ImportError:
+    # Try alternative import path
+    import sys
+    import os
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    sys.path.insert(0, os.path.join(current_dir, "lm-harness-integration"))
+    from only_benchmarks import CORE_BENCHMARKS
+
+from .contrastive_pairs import ContrastivePairSet
+from .managed_cached_benchmarks import ManagedCachedBenchmarks, get_managed_cache
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class BenchmarkSample:
+    """A single sample from a benchmark."""
+    benchmark_name: str
+    sample_data: Dict[str, Any]
+    tags: List[str]
+
+
+class MixedBenchmarkSampler:
+    """
+    Samples randomly from multiple benchmarks based on tags.
+    
+    This creates more robust classifiers by training on diverse data
+    from multiple sources rather than a single benchmark.
+    """
+    
+    def __init__(self, cache_dir: str = "./benchmark_cache"):
+        """
+        Initialize the mixed benchmark sampler.
+        
+        Args:
+            cache_dir: Directory for cached benchmark data
+        """
+        self.cache_dir = cache_dir
+        self.managed_cache = get_managed_cache(cache_dir)
+        self._benchmark_registry = self._build_benchmark_registry()
+    
+    def _build_benchmark_registry(self) -> Dict[str, List[str]]:
+        """Build a registry mapping tags to benchmark names."""
+        tag_to_benchmarks = defaultdict(list)
+        
+        for benchmark_name, config in CORE_BENCHMARKS.items():
+            tags = config.get("tags", [])
+            for tag in tags:
+                tag_to_benchmarks[tag].append(benchmark_name)
+        
+        return dict(tag_to_benchmarks)
+    
+    def get_benchmarks_by_tag(self, tag: str) -> List[str]:
+        """Get all benchmarks that have a specific tag."""
+        return self._benchmark_registry.get(tag, [])
+    
+    def get_benchmarks_by_tags(self, tags: List[str], mode: str = "any") -> List[str]:
+        """
+        Get benchmarks that match the given tags.
+        
+        Args:
+            tags: List of tags to match
+            mode: "any" (benchmark has at least one tag) or "all" (benchmark has all tags)
+            
+        Returns:
+            List of benchmark names matching the criteria
+        """
+        if mode == "any":
+            # Get benchmarks that have ANY of the specified tags
+            matching_benchmarks = set()
+            for tag in tags:
+                matching_benchmarks.update(self.get_benchmarks_by_tag(tag))
+            return list(matching_benchmarks)
+        
+        elif mode == "all":
+            # Get benchmarks that have ALL of the specified tags
+            if not tags:
+                return []
+            
+            # Start with benchmarks that have the first tag
+            matching_benchmarks = set(self.get_benchmarks_by_tag(tags[0]))
+            
+            # Intersect with benchmarks for each additional tag
+            for tag in tags[1:]:
+                matching_benchmarks &= set(self.get_benchmarks_by_tag(tag))
+            
+            return list(matching_benchmarks)
+        
+        else:
+            raise ValueError(f"Invalid mode: {mode}. Use 'any' or 'all'")
+    
+    def sample_mixed_dataset(
+        self,
+        tags: List[str],
+        total_samples: int,
+        split_ratio: float = 0.8,
+        random_seed: Optional[int] = None,
+        tag_mode: str = "any",
+        benchmark_weights: Optional[Dict[str, float]] = None
+    ) -> Tuple[List[BenchmarkSample], List[BenchmarkSample]]:
+        """
+        Sample a mixed dataset from benchmarks matching the given tags.
+        
+        Args:
+            tags: Tags to filter benchmarks (e.g., ["coding", "python"])
+            total_samples: Total number of samples to collect
+            split_ratio: Train/test split ratio
+            random_seed: Random seed for reproducibility
+            tag_mode: "any" or "all" for tag matching
+            benchmark_weights: Optional weights for sampling probability per benchmark
+            
+        Returns:
+            Tuple of (train_samples, test_samples)
+        """
+        if random_seed is not None:
+            random.seed(random_seed)
+        
+        # Get matching benchmarks
+        matching_benchmarks = self.get_benchmarks_by_tags(tags, mode=tag_mode)
+        
+        if not matching_benchmarks:
+            raise ValueError(f"No benchmarks found with tags {tags} (mode={tag_mode})")
+        
+        logger.info(f"Found {len(matching_benchmarks)} benchmarks matching tags {tags}")
+        logger.info(f"Matching benchmarks: {matching_benchmarks[:10]}...")  # Show first 10
+        
+        # Collect all available samples from matching benchmarks
+        all_samples = []
+        benchmark_sample_counts = {}
+        
+        # Skip benchmarks that require code execution permission
+        code_execution_benchmarks = {"apps", "ds1000", "mercury"}
+        
+        for benchmark_name in matching_benchmarks:
+            # Skip benchmarks that require code execution for safety
+            if benchmark_name in code_execution_benchmarks:
+                logger.info(f"Skipping {benchmark_name} (requires code execution permission)")
+                continue
+                
+            try:
+                # Get samples from this benchmark
+                samples_per_benchmark = max(10, total_samples // len(matching_benchmarks))
+                
+                cached_samples = self.managed_cache.get_task_samples(
+                    task_name=benchmark_name,
+                    limit=samples_per_benchmark,
+                    force_fresh=False
+                )
+                
+                # Convert to BenchmarkSample objects
+                for sample in cached_samples:
+                    benchmark_sample = BenchmarkSample(
+                        benchmark_name=benchmark_name,
+                        sample_data=sample,
+                        tags=CORE_BENCHMARKS[benchmark_name].get("tags", [])
+                    )
+                    all_samples.append(benchmark_sample)
+                
+                benchmark_sample_counts[benchmark_name] = len(cached_samples)
+                
+            except Exception as e:
+                logger.warning(f"Failed to load samples from {benchmark_name}: {e}")
+                continue
+        
+        if not all_samples:
+            raise ValueError(f"No samples could be loaded from any benchmark with tags {tags}")
+        
+        logger.info(f"Collected {len(all_samples)} total samples from {len(benchmark_sample_counts)} benchmarks")
+        for benchmark, count in benchmark_sample_counts.items():
+            logger.debug(f"  {benchmark}: {count} samples")
+        
+        # Apply benchmark weights if provided
+        if benchmark_weights:
+            weighted_samples = []
+            for sample in all_samples:
+                weight = benchmark_weights.get(sample.benchmark_name, 1.0)
+                # Duplicate samples based on weight (simple approach)
+                weighted_samples.extend([sample] * int(weight))
+            all_samples = weighted_samples
+        
+        # Randomly sample and shuffle
+        if len(all_samples) > total_samples:
+            all_samples = random.sample(all_samples, total_samples)
+        else:
+            # If we have fewer samples than requested, use all and log warning
+            logger.warning(f"Only {len(all_samples)} samples available, requested {total_samples}")
+        
+        random.shuffle(all_samples)
+        
+        # Split into train/test
+        split_point = int(len(all_samples) * split_ratio)
+        train_samples = all_samples[:split_point]
+        test_samples = all_samples[split_point:]
+        
+        # Log distribution
+        train_dist = defaultdict(int)
+        test_dist = defaultdict(int)
+        
+        for sample in train_samples:
+            train_dist[sample.benchmark_name] += 1
+        
+        for sample in test_samples:
+            test_dist[sample.benchmark_name] += 1
+        
+        logger.info(f"Train set: {len(train_samples)} samples from {len(train_dist)} benchmarks")
+        logger.info(f"Test set: {len(test_samples)} samples from {len(test_dist)} benchmarks")
+        
+        return train_samples, test_samples
+    
+    def extract_contrastive_pairs_from_mixed_samples(
+        self,
+        samples: List[BenchmarkSample]
+    ) -> List[Dict[str, Any]]:
+        """
+        Extract contrastive pairs from mixed benchmark samples.
+        
+        Args:
+            samples: List of BenchmarkSample objects
+            
+        Returns:
+            List of contrastive pairs with question, correct_answer, incorrect_answer
+        """
+        contrastive_pairs = []
+        
+        for sample in samples:
+            try:
+                # Each sample already has normalized QA pair from managed cache
+                qa_pair = sample.sample_data.get("normalized", {})
+                
+                if qa_pair and all(k in qa_pair for k in ["question", "correct_answer", "incorrect_answer"]):
+                    # Add benchmark source info
+                    qa_pair["source_benchmark"] = sample.benchmark_name
+                    qa_pair["tags"] = sample.tags
+                    contrastive_pairs.append(qa_pair)
+                else:
+                    logger.warning(f"Invalid QA pair from {sample.benchmark_name}")
+                    
+            except Exception as e:
+                logger.warning(f"Failed to extract pair from {sample.benchmark_name}: {e}")
+                continue
+        
+        logger.info(f"Extracted {len(contrastive_pairs)} contrastive pairs from mixed samples")
+        
+        return contrastive_pairs
+    
+    def create_mixed_contrastive_pair_set(
+        self,
+        tags: List[str],
+        total_samples: int,
+        name: Optional[str] = None,
+        **kwargs
+    ) -> ContrastivePairSet:
+        """
+        Create a ContrastivePairSet from mixed benchmark samples.
+        
+        Args:
+            tags: Tags to filter benchmarks
+            total_samples: Number of samples to include
+            name: Name for the pair set (auto-generated if not provided)
+            **kwargs: Additional arguments for sample_mixed_dataset
+            
+        Returns:
+            ContrastivePairSet ready for training
+        """
+        # Sample mixed dataset
+        train_samples, test_samples = self.sample_mixed_dataset(
+            tags=tags,
+            total_samples=total_samples,
+            **kwargs
+        )
+        
+        # Extract contrastive pairs
+        all_samples = train_samples + test_samples
+        contrastive_pairs = self.extract_contrastive_pairs_from_mixed_samples(all_samples)
+        
+        # Create name if not provided
+        if name is None:
+            name = f"mixed_{'_'.join(tags)}_{total_samples}_samples"
+        
+        # Create ContrastivePairSet
+        return ContrastivePairSet.from_contrastive_pairs(
+            name=name,
+            contrastive_pairs=contrastive_pairs,
+            task_type="mixed_benchmark"
+        )
+
+
+def sample_benchmarks_by_tag(
+    tag: str,
+    samples_per_benchmark: int = 10,
+    max_benchmarks: Optional[int] = None,
+    random_seed: Optional[int] = None
+) -> Dict[str, List[Dict[str, Any]]]:
+    """
+    Convenience function to sample from all benchmarks with a specific tag.
+    
+    Args:
+        tag: Tag to filter benchmarks (e.g., "coding")
+        samples_per_benchmark: Number of samples from each benchmark
+        max_benchmarks: Maximum number of benchmarks to sample from
+        random_seed: Random seed for reproducibility
+        
+    Returns:
+        Dictionary mapping benchmark names to their samples
+    """
+    sampler = MixedBenchmarkSampler()
+    
+    # Get all benchmarks with the tag
+    benchmarks = sampler.get_benchmarks_by_tag(tag)
+    
+    if max_benchmarks and len(benchmarks) > max_benchmarks:
+        if random_seed is not None:
+            random.seed(random_seed)
+        benchmarks = random.sample(benchmarks, max_benchmarks)
+    
+    # Sample from each benchmark
+    results = {}
+    cache = get_managed_cache()
+    
+    for benchmark_name in benchmarks:
+        try:
+            samples = cache.get_task_samples(
+                task_name=benchmark_name,
+                limit=samples_per_benchmark,
+                force_fresh=False
+            )
+            results[benchmark_name] = samples
+            logger.info(f"Sampled {len(samples)} from {benchmark_name}")
+            
+        except Exception as e:
+            logger.warning(f"Failed to sample from {benchmark_name}: {e}")
+            continue
+    
+    return results

wisent/core/model_config_manager.py

@@ -0,0 +1,330 @@
+"""
+Model Configuration Manager for storing and retrieving optimal parameters per model.
+"""
+
+import os
+import json
+import logging
+from typing import Dict, Any, Optional, List
+from datetime import datetime
+import hashlib
+import numpy as np
+
+
+class NumpyEncoder(json.JSONEncoder):
+    """Custom JSON encoder to handle numpy types."""
+    def default(self, obj):
+        if isinstance(obj, (np.integer, np.int64)):
+            return int(obj)
+        if isinstance(obj, (np.floating, np.float64)):
+            return float(obj)
+        if isinstance(obj, np.ndarray):
+            return obj.tolist()
+        return super().default(obj)
+
+logger = logging.getLogger(__name__)
+
+
+class ModelConfigManager:
+    """Manages model-specific configuration files for optimal parameters."""
+    
+    def __init__(self, config_dir: Optional[str] = None):
+        """
+        Initialize the ModelConfigManager.
+        
+        Args:
+            config_dir: Directory to store config files. If None, uses default location.
+        """
+        if config_dir is None:
+            # Use ~/.wisent-guard/model_configs/ as default
+            home_dir = os.path.expanduser("~")
+            self.config_dir = os.path.join(home_dir, ".wisent-guard", "model_configs")
+        else:
+            self.config_dir = config_dir
+            
+        # Create directory if it doesn't exist
+        os.makedirs(self.config_dir, exist_ok=True)
+        
+    def _sanitize_model_name(self, model_name: str) -> str:
+        """
+        Convert model name to a safe filename.
+        
+        Args:
+            model_name: Original model name (e.g., "meta-llama/Llama-3.1-8B-Instruct")
+            
+        Returns:
+            Sanitized filename (e.g., "meta-llama_Llama-3.1-8B-Instruct")
+        """
+        # Replace problematic characters
+        sanitized = model_name.replace("/", "_").replace("\\", "_").replace(":", "_")
+        # Remove any other problematic characters
+        sanitized = "".join(c for c in sanitized if c.isalnum() or c in "._-")
+        return sanitized
+    
+    def _get_config_path(self, model_name: str) -> str:
+        """Get the full path to the config file for a model."""
+        sanitized_name = self._sanitize_model_name(model_name)
+        return os.path.join(self.config_dir, f"{sanitized_name}.json")
+    
+    def save_model_config(
+        self,
+        model_name: str,
+        classification_layer: int,
+        steering_layer: Optional[int] = None,
+        token_aggregation: str = "average",
+        detection_threshold: float = 0.6,
+        optimization_method: str = "manual",
+        optimization_metrics: Optional[Dict[str, Any]] = None,
+        task_specific_overrides: Optional[Dict[str, Dict[str, Any]]] = None
+    ) -> str:
+        """
+        Save optimal parameters for a model.
+        
+        Args:
+            model_name: Name/path of the model
+            classification_layer: Optimal layer for classification
+            steering_layer: Optimal layer for steering (defaults to classification_layer)
+            token_aggregation: Token aggregation method
+            detection_threshold: Detection threshold
+            optimization_method: How these parameters were determined
+            optimization_metrics: Metrics from optimization process
+            task_specific_overrides: Task-specific parameter overrides
+            
+        Returns:
+            Path to the saved config file
+        """
+        if steering_layer is None:
+            steering_layer = classification_layer
+            
+        config_data = {
+            "model_name": model_name,
+            "created_date": datetime.now().isoformat(),
+            "optimization_method": optimization_method,
+            "optimal_parameters": {
+                "classification_layer": classification_layer,
+                "steering_layer": steering_layer,
+                "token_aggregation": token_aggregation,
+                "detection_threshold": detection_threshold
+            },
+            "task_specific_overrides": task_specific_overrides or {},
+            "optimization_metrics": optimization_metrics or {},
+            "config_version": "1.0"
+        }
+        
+        config_path = self._get_config_path(model_name)
+        
+        try:
+            with open(config_path, 'w') as f:
+                json.dump(config_data, f, indent=2, cls=NumpyEncoder)
+            
+            logger.info(f"✅ Model configuration saved: {config_path}")
+            logger.info(f"   • Classification layer: {classification_layer}")
+            logger.info(f"   • Steering layer: {steering_layer}")
+            logger.info(f"   • Token aggregation: {token_aggregation}")
+            logger.info(f"   • Detection threshold: {detection_threshold}")
+            
+            return config_path
+            
+        except Exception as e:
+            logger.error(f"❌ Failed to save model configuration: {e}")
+            raise
+    
+    def load_model_config(self, model_name: str) -> Optional[Dict[str, Any]]:
+        """
+        Load optimal parameters for a model.
+        
+        Args:
+            model_name: Name/path of the model
+            
+        Returns:
+            Configuration dictionary if found, None otherwise
+        """
+        config_path = self._get_config_path(model_name)
+        
+        if not os.path.exists(config_path):
+            return None
+            
+        try:
+            with open(config_path, 'r') as f:
+                config_data = json.load(f)
+            
+            logger.debug(f"📄 Loaded model configuration: {config_path}")
+            return config_data
+            
+        except Exception as e:
+            logger.warning(f"⚠️ Failed to load model configuration: {e}")
+            return None
+    
+    def has_model_config(self, model_name: str) -> bool:
+        """Check if a model has a saved configuration."""
+        config_path = self._get_config_path(model_name)
+        return os.path.exists(config_path)
+    
+    def update_model_config(self, model_name: str, config_data: Dict[str, Any]) -> str:
+        """
+        Update an existing model configuration.
+        
+        Args:
+            model_name: Name/path of the model
+            config_data: Updated configuration dictionary
+            
+        Returns:
+            Path to the saved config file
+        """
+        config_path = self._get_config_path(model_name)
+        
+        # Update timestamp
+        config_data["updated_date"] = datetime.now().isoformat()
+        
+        try:
+            with open(config_path, 'w') as f:
+                json.dump(config_data, f, indent=2, cls=NumpyEncoder)
+            
+            logger.info(f"✅ Model configuration updated: {config_path}")
+            return config_path
+            
+        except Exception as e:
+            logger.error(f"❌ Failed to update model configuration: {e}")
+            raise
+    
+    def get_optimal_parameters(
+        self, 
+        model_name: str, 
+        task_name: Optional[str] = None
+    ) -> Optional[Dict[str, Any]]:
+        """
+        Get optimal parameters for a model, with optional task-specific overrides.
+        
+        Args:
+            model_name: Name/path of the model
+            task_name: Specific task name for overrides
+            
+        Returns:
+            Dictionary of optimal parameters or None if no config exists
+        """
+        config = self.load_model_config(model_name)
+        if not config:
+            return None
+        
+        # Start with base optimal parameters
+        optimal_params = config.get("optimal_parameters", {}).copy()
+        
+        # Apply task-specific overrides if available
+        if task_name and "task_specific_overrides" in config:
+            task_overrides = config["task_specific_overrides"].get(task_name, {})
+            optimal_params.update(task_overrides)
+        
+        return optimal_params
+    
+    def get_optimal_sample_size(
+        self,
+        model_name: str,
+        task_name: str,
+        layer: int
+    ) -> Optional[int]:
+        """
+        Get optimal sample size for a specific task and layer.
+        
+        Args:
+            model_name: Name/path of the model
+            task_name: Task name
+            layer: Layer index
+            
+        Returns:
+            Optimal sample size or None if not found
+        """
+        config = self.load_model_config(model_name)
+        if not config:
+            return None
+            
+        # Check if optimal_sample_sizes exists
+        if "optimal_sample_sizes" not in config:
+            return None
+            
+        # Navigate the nested structure: optimal_sample_sizes[task][layer]
+        task_sizes = config["optimal_sample_sizes"].get(task_name, {})
+        sample_size = task_sizes.get(str(layer), None)
+        
+        return sample_size
+    
+    def list_model_configs(self) -> List[Dict[str, Any]]:
+        """
+        List all available model configurations.
+        
+        Returns:
+            List of configuration summaries
+        """
+        configs = []
+        
+        if not os.path.exists(self.config_dir):
+            return configs
+        
+        for filename in os.listdir(self.config_dir):
+            if filename.endswith('.json'):
+                try:
+                    config_path = os.path.join(self.config_dir, filename)
+                    with open(config_path, 'r') as f:
+                        config_data = json.load(f)
+                    
+                    summary = {
+                        "model_name": config_data.get("model_name", "unknown"),
+                        "created_date": config_data.get("created_date", "unknown"),
+                        "optimization_method": config_data.get("optimization_method", "unknown"),
+                        "classification_layer": config_data.get("optimal_parameters", {}).get("classification_layer"),
+                        "steering_layer": config_data.get("optimal_parameters", {}).get("steering_layer"),
+                        "config_file": filename
+                    }
+                    configs.append(summary)
+                    
+                except Exception as e:
+                    logger.warning(f"⚠️ Failed to read config file {filename}: {e}")
+        
+        return configs
+    
+    def remove_model_config(self, model_name: str) -> bool:
+        """
+        Remove a model configuration.
+        
+        Args:
+            model_name: Name/path of the model
+            
+        Returns:
+            True if removed successfully, False otherwise
+        """
+        config_path = self._get_config_path(model_name)
+        
+        if not os.path.exists(config_path):
+            logger.warning(f"⚠️ No configuration found for model: {model_name}")
+            return False
+        
+        try:
+            os.remove(config_path)
+            logger.info(f"✅ Removed model configuration: {config_path}")
+            return True
+            
+        except Exception as e:
+            logger.error(f"❌ Failed to remove model configuration: {e}")
+            return False
+
+
+# Convenience functions for easy access
+_default_manager = None
+
+def get_default_manager() -> ModelConfigManager:
+    """Get the default ModelConfigManager instance."""
+    global _default_manager
+    if _default_manager is None:
+        _default_manager = ModelConfigManager()
+    return _default_manager
+
+def save_model_config(model_name: str, **kwargs) -> str:
+    """Save model configuration using default manager."""
+    return get_default_manager().save_model_config(model_name, **kwargs)
+
+def load_model_config(model_name: str) -> Optional[Dict[str, Any]]:
+    """Load model configuration using default manager."""
+    return get_default_manager().load_model_config(model_name)
+
+def get_optimal_parameters(model_name: str, task_name: Optional[str] = None) -> Optional[Dict[str, Any]]:
+    """Get optimal parameters using default manager."""
+    return get_default_manager().get_optimal_parameters(model_name, task_name)