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

wisent/core/optuna/steering/bigcode_evaluator_wrapper.py

@@ -0,0 +1,188 @@
+"""
+BigCode evaluator wrapper for optuna pipeline integration.
+
+This module provides a clean interface for integrating BigCode code execution
+evaluation with the optuna optimization pipeline.
+"""
+
+import logging
+from typing import Any, Dict, List, Tuple
+
+from wisent.core.bigcode_extractors import get_bigcode_extractor
+from wisent.core.bigcode_integration import BigCodeEvaluator, is_bigcode_task
+from wisent.parameters.task_config import CODING_TASKS
+
+logger = logging.getLogger(__name__)
+
+
+class OptunaBigCodeEvaluator:
+    """
+    Wrapper for BigCode evaluation in optuna pipeline.
+
+    This class provides a clean interface for evaluating coding tasks
+    using actual code execution instead of string comparison.
+    """
+
+    def __init__(self, docker_executor=None):
+        """
+        Initialize the evaluator.
+
+        Args:
+            docker_executor: Optional Docker executor for secure code execution
+        """
+        self.bigcode_evaluator = BigCodeEvaluator(docker_executor)
+        self.code_extractor = None  # Will be set per task
+
+    def is_coding_task(self, task_name: str) -> bool:
+        """
+        Check if a task requires code execution evaluation.
+
+        Args:
+            task_name: Name of the task
+
+        Returns:
+            True if task requires code execution, False otherwise
+        """
+        if not task_name:
+            return False
+        return task_name.lower() in CODING_TASKS or is_bigcode_task(task_name)
+
+    def evaluate_predictions(
+        self, predictions: List[str], task_docs: List[Dict[str, Any]], task_name: str
+    ) -> List[Dict[str, Any]]:
+        """
+        Evaluate model predictions using code execution.
+
+        Args:
+            predictions: List of model-generated code predictions
+            task_docs: List of original task documents with test cases
+            task_name: Name of the task
+
+        Returns:
+            List of evaluation results for each prediction
+        """
+        if not self.is_coding_task(task_name):
+            raise ValueError(f"Task {task_name} is not a coding task")
+
+        if len(predictions) != len(task_docs):
+            raise ValueError(f"Mismatch: {len(predictions)} predictions vs {len(task_docs)} task docs")
+
+        results = []
+
+        for i, (prediction, task_doc) in enumerate(zip(predictions, task_docs)):
+            try:
+                # Get the appropriate extractor for this task
+                code_extractor = get_bigcode_extractor(task_name)
+
+                # Extract code from the prediction
+                extracted_code = code_extractor.extract_code_from_answer(prediction)
+
+                if not extracted_code.strip():
+                    logger.warning(f"No code extracted from prediction {i}: {prediction[:100]}...")
+                    result = {
+                        "passed": False,
+                        "error": "No code extracted from prediction",
+                        "extracted_code": extracted_code,
+                        "original_prediction": prediction,
+                    }
+                else:
+                    # Execute the code against test cases
+                    result = self.bigcode_evaluator._execute_and_test(task_doc, extracted_code, task_name)
+                    result["extracted_code"] = extracted_code
+                    result["original_prediction"] = prediction
+
+                results.append(result)
+
+            except Exception as e:
+                logger.warning(f"Error evaluating prediction {i}: {e}")
+                results.append(
+                    {"passed": False, "error": str(e), "extracted_code": "", "original_prediction": prediction}
+                )
+
+        return results
+
+    def calculate_accuracy(self, evaluation_results: List[Dict[str, Any]]) -> Dict[str, Any]:
+        """
+        Calculate accuracy metrics from evaluation results.
+
+        Args:
+            evaluation_results: List of evaluation results from evaluate_predictions
+
+        Returns:
+            Dictionary with accuracy metrics
+        """
+        total_samples = len(evaluation_results)
+        if total_samples == 0:
+            return {
+                "accuracy": 0.0,
+                "total_samples": 0,
+                "correct": 0,
+                "incorrect": 0,
+                "pass_count": 0,
+                "fail_count": 0,
+                "error_count": 0,
+            }
+
+        pass_count = sum(1 for result in evaluation_results if result.get("passed", False))
+        fail_count = sum(
+            1 for result in evaluation_results if result.get("passed", False) == False and not result.get("error")
+        )
+        error_count = sum(1 for result in evaluation_results if result.get("error"))
+
+        accuracy = pass_count / total_samples
+
+        return {
+            "accuracy": accuracy,
+            "total_samples": total_samples,
+            "correct": pass_count,
+            "incorrect": fail_count + error_count,
+            "pass_count": pass_count,
+            "fail_count": fail_count,
+            "error_count": error_count,
+            "evaluation_method": "bigcode_execution",
+        }
+
+    def evaluate_and_calculate_accuracy(
+        self, predictions: List[str], task_docs: List[Dict[str, Any]], task_name: str
+    ) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]:
+        """
+        Convenience method to evaluate predictions and calculate accuracy.
+
+        Args:
+            predictions: List of model-generated code predictions
+            task_docs: List of original task documents with test cases
+            task_name: Name of the task
+
+        Returns:
+            Tuple of (evaluation_results, accuracy_metrics)
+        """
+        evaluation_results = self.evaluate_predictions(predictions, task_docs, task_name)
+        accuracy_metrics = self.calculate_accuracy(evaluation_results)
+
+        logger.info(
+            f"BigCode evaluation for {task_name}: "
+            f"{accuracy_metrics['pass_count']}/{accuracy_metrics['total_samples']} passed "
+            f"({accuracy_metrics['accuracy']:.3f} accuracy)"
+        )
+
+        return evaluation_results, accuracy_metrics
+
+
+# Global instance for easy access
+_optuna_bigcode_evaluator = None
+
+
+def get_optuna_bigcode_evaluator(docker_executor=None) -> OptunaBigCodeEvaluator:
+    """
+    Get global OptunaBigCodeEvaluator instance.
+
+    Args:
+        docker_executor: Optional Docker executor
+
+    Returns:
+        OptunaBigCodeEvaluator instance
+    """
+    global _optuna_bigcode_evaluator
+    if _optuna_bigcode_evaluator is None:
+        _optuna_bigcode_evaluator = OptunaBigCodeEvaluator(docker_executor)
+    return _optuna_bigcode_evaluator

wisent/core/optuna/steering/data_utils.py

@@ -0,0 +1,342 @@
+"""
+Data loading and processing utilities for comprehensive evaluation.
+"""
+
+import logging
+from typing import Dict, List, Tuple
+
+import numpy as np
+import torch
+from tqdm import tqdm
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+# Import LMEvalHarnessGroundTruth for intelligent evaluation (same approach as CLI)
+from wisent.core.lm_eval_harness_ground_truth import LMEvalHarnessGroundTruth
+
+# Import task interface for dynamic task loading
+from wisent.core.task_interface import get_task
+from wisent.core.utils.device import empty_device_cache, preferred_dtype, resolve_default_device
+
+logger = logging.getLogger(__name__)
+
+
+def load_dataset_samples(dataset_name: str, limit: int) -> List[Dict]:
+    """Load samples from a dataset using the unified task interface."""
+    logger.info(f"Loading {limit} samples from {dataset_name}...")
+
+    try:
+        # Use the unified task interface to get any registered task
+        task = get_task(dataset_name, limit=limit)
+        samples = task.load_data(limit=limit)
+
+        logger.info(f"Loaded {len(samples)} samples from {dataset_name} via {task.__class__.__name__}")
+        return samples
+
+    except Exception as e:
+        logger.error(f"Failed to load {dataset_name}: {e}")
+        # Provide helpful error message with available tasks
+        try:
+            from ...task_interface import list_tasks
+
+            available_tasks = list_tasks()
+            logger.error(f"Available tasks: {available_tasks}")
+        except:
+            pass
+        raise
+
+
+def extract_activations_with_hook(
+    model, tokenizer, texts: List[str], layer: int, batch_size: int, max_length: int, device: torch.device
+) -> np.ndarray:
+    """Extract activations from a specific layer using hooks."""
+    activations = []
+
+    def hook_fn(module, input, output):
+        # Handle different output formats (some layers return tuples)
+        if isinstance(output, tuple):
+            hidden_states = output[0]  # First element is usually hidden states
+        else:
+            hidden_states = output
+
+        # Extract last token activations (typical for causal LM)
+        if len(hidden_states.shape) == 3:  # [batch, seq, hidden]
+            last_token_acts = hidden_states[:, -1, :].detach().cpu().numpy()
+            activations.extend(last_token_acts)
+
+    # Register hook
+    if hasattr(model, "transformer"):  # GPT-style models
+        target_layer = model.transformer.h[layer]
+    elif hasattr(model, "model"):  # Some other architectures
+        target_layer = model.model.layers[layer]
+    else:
+        raise ValueError("Unknown model architecture")
+
+    handle = target_layer.register_forward_hook(hook_fn)
+
+    try:
+        # Process texts in batches
+        for i in tqdm(range(0, len(texts), batch_size), desc=f"Extracting activations (layer {layer})"):
+            batch_texts = texts[i : i + batch_size]
+
+            inputs = tokenizer(
+                batch_texts, return_tensors="pt", padding=True, truncation=True, max_length=max_length
+            ).to(device)
+
+            with torch.no_grad():
+                _ = model(**inputs)
+
+    finally:
+        handle.remove()
+
+    return np.array(activations)
+
+
+def generate_benchmark_predictions(
+    model,
+    tokenizer,
+    samples: List[Dict],
+    batch_size: int,
+    max_length: int,
+    device: torch.device,
+    task_name: str,
+    max_new_tokens: int,
+    preserve_task_docs: bool = False,
+) -> Tuple[List[str], List[str], List[Dict]]:
+    """Generate model predictions for benchmark evaluation using task extractor with batching.
+
+    Args:
+        preserve_task_docs: If True, returns original task documents alongside predictions
+
+    Returns:
+        Tuple of (predictions, ground_truths, task_docs) if preserve_task_docs=True
+        Tuple of (predictions, ground_truths, []) if preserve_task_docs=False
+    """
+    predictions = []
+    ground_truths = []
+    task_docs = [] if preserve_task_docs else []
+
+    # Get the task and its extractor
+    task = get_task(task_name)
+    extractor = task.get_extractor()
+
+    # First, extract all questions and answers
+    questions = []
+    answers = []
+
+    valid_samples = []  # Keep track of samples that produce valid QA pairs
+
+    for sample in samples:
+        qa_pair = extractor.extract_qa_pair(sample, task)
+        if not qa_pair:
+            logger.warning(f"Skipping sample - extractor couldn't extract QA pair: {sample.keys()}")
+            continue
+        questions.append(qa_pair["formatted_question"])
+        answers.append(qa_pair["correct_answer"])
+
+        if preserve_task_docs:
+            valid_samples.append(sample)
+
+    # Process in batches
+    for i in tqdm(range(0, len(questions), batch_size), desc="Generating benchmark predictions"):
+        batch_questions = questions[i : i + batch_size]
+        batch_answers = answers[i : i + batch_size]
+
+        # Tokenize batch
+        inputs = tokenizer(
+            batch_questions, return_tensors="pt", padding=True, truncation=True, max_length=max_length
+        ).to(device)
+
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs, max_new_tokens=max_new_tokens, do_sample=True, pad_token_id=tokenizer.eos_token_id
+            )
+
+        # Extract generated text for each item in batch
+        for j, output in enumerate(outputs):
+            input_length = inputs["input_ids"][j].shape[0]
+            generated = tokenizer.decode(output[input_length:], skip_special_tokens=True)
+            generated = generated.strip()
+            predictions.append(generated)
+
+        # Add ground truths
+        ground_truths.extend(batch_answers)
+
+    # Add task docs if requested
+    if preserve_task_docs:
+        task_docs = valid_samples[: len(predictions)]  # Ensure same length as predictions
+
+    return predictions, ground_truths, task_docs
+
+
+def create_probe_training_data(
+    model,
+    tokenizer,
+    samples: List[Dict],
+    layer: int,
+    batch_size: int,
+    max_length: int,
+    device: torch.device,
+    task_name: str,
+    max_new_tokens: int = 200,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Create training data for probes: activations -> correctness labels using task extractor with batched generation."""
+    texts = []
+    labels = []
+
+    # Get the task and its extractor
+    task = get_task(task_name)
+    extractor = task.get_extractor()
+
+    # Pre-extract all questions and answers for batched generation
+    questions = []
+    correct_answers = []
+
+    for sample in samples:
+        qa_pair = extractor.extract_qa_pair(sample, task)
+        if not qa_pair:
+            continue
+        questions.append(qa_pair["formatted_question"])
+        correct_answers.append(qa_pair["correct_answer"])
+
+    # Generate predictions in batches
+    generated_answers = []
+
+    for i in tqdm(range(0, len(questions), batch_size), desc=f"Generating probe data (layer {layer})"):
+        batch_questions = questions[i : i + batch_size]
+
+        # Tokenize batch
+        inputs = tokenizer(
+            batch_questions, return_tensors="pt", padding=True, truncation=True, max_length=max_length
+        ).to(device)
+
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs, max_new_tokens=max_new_tokens, do_sample=True, pad_token_id=tokenizer.eos_token_id
+            )
+
+        # Extract generated text for each item in batch
+        for j, output in enumerate(outputs):
+            input_length = inputs["input_ids"][j].shape[0]
+            generated = tokenizer.decode(output[input_length:], skip_special_tokens=True)
+            generated = generated.strip()
+            generated_answers.append(generated)
+
+    # Now process each question-answer pair for probe training data
+    evaluator = LMEvalHarnessGroundTruth(task_name)
+
+    for question, correct_answer, generated in zip(questions, correct_answers, generated_answers):
+        # Create examples with model's actual prediction
+        correct_text = f"{question} {correct_answer}"
+        incorrect_text = f"{question} {generated}"
+
+        texts.extend([correct_text, incorrect_text])
+
+        # Evaluate if prediction is correct using LMEvalHarnessGroundTruth
+        try:
+            # Create response data format expected by _evaluate_with_lm_eval_metrics
+            response_data = [
+                {
+                    "generated_response": generated,
+                    "ground_truth": correct_answer,
+                    "question": "evaluation_question",  # Required field for evaluation
+                }
+            ]
+
+            # Use the same evaluation logic as CLI
+            eval_results = evaluator._evaluate_with_lm_eval_metrics(task_name, response_data, None)
+
+            # Extract the result - accuracy > 0 means at least one correct
+            is_correct = eval_results.get("accuracy", 0.0) > 0.0
+
+        except Exception as e:
+            logger.warning(f"LMEvalHarnessGroundTruth failed, using exact match fallback: {e}")
+            is_correct = generated.strip().lower() == correct_answer.strip().lower()
+
+        labels.extend([1, 1 if is_correct else 0])
+
+    # Extract activations
+    activations = extract_activations_with_hook(model, tokenizer, texts, layer, batch_size, max_length, device)
+
+    return activations, np.array(labels)
+
+
+def load_model_and_tokenizer(model_name: str, device: torch.device):
+    """Load model and tokenizer with proper configuration."""
+    logger.info(f"Loading model {model_name} (ONCE)...")
+    device_kind = device.type
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+    # Configure tokenizer for decoder-only models
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    # Set left padding for decoder-only models (required for correct generation)
+    tokenizer.padding_side = "left"
+
+    torch_dtype = preferred_dtype(device_kind)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        torch_dtype=torch_dtype,
+        low_cpu_mem_usage=True,
+    )
+    model.to(device)
+    model.eval()
+
+    # Log memory usage
+    if device_kind == "cuda" and torch.cuda.is_available():
+        memory_gb = torch.cuda.memory_allocated() / 1024**3
+        logger.info(f"✓ Model loaded on {device}, GPU memory: {memory_gb:.2f} GB")
+    elif device_kind == "mps" and hasattr(torch, "mps"):
+        try:
+            memory_gb = torch.mps.current_allocated_memory() / 1024**3
+            logger.info(f"✓ Model loaded on {device}, MPS memory: {memory_gb:.2f} GB")
+        except AttributeError:
+            logger.info(f"✓ Model loaded on {device}")
+
+    return model, tokenizer
+
+
+def free_model_memory(model, tokenizer):
+    """Free model memory after activation extraction."""
+    logger.info("🧹 Freeing model memory...")
+    device_kind = None
+    if hasattr(model, "parameters"):
+        try:
+            device_kind = next(model.parameters()).device.type
+        except StopIteration:
+            pass
+    del model
+    del tokenizer
+    import gc
+
+    gc.collect()
+    kind_for_cleanup = device_kind or resolve_default_device()
+    empty_device_cache(kind_for_cleanup)
+    if kind_for_cleanup == "cuda" and torch.cuda.is_available():
+        memory_gb = torch.cuda.memory_allocated() / 1024**3
+        logger.info(f"GPU memory after cleanup: {memory_gb:.2f} GB")
+    elif kind_for_cleanup == "mps" and hasattr(torch, "mps"):
+        try:
+            memory_gb = torch.mps.current_allocated_memory() / 1024**3
+            logger.info(f"MPS memory after cleanup: {memory_gb:.2f} GB")
+        except AttributeError:
+            pass
+
+
+def get_task_contrastive_pairs(samples: List[Dict], task_name: str) -> List[Dict]:
+    """Extract contrastive pairs from samples using the task's extractor."""
+    contrastive_pairs = []
+
+    # Get the task and its extractor
+    task = get_task(task_name)
+    extractor = task.get_extractor()
+
+    for sample in samples:
+        # Use the task's extractor to get contrastive pair
+        pair = extractor.extract_contrastive_pair(sample, task)
+        if pair:
+            contrastive_pairs.append(pair)
+
+    logger.info(f"Extracted {len(contrastive_pairs)} contrastive pairs from {len(samples)} samples")
+    return contrastive_pairs