hackagent 0.1.0__py3-none-any.whl

hackagent/attacks/AdvPrefix/step7_evaluate_responses.py

@@ -0,0 +1,289 @@
+import logging
+import pandas as pd
+import multiprocessing
+import os
+from typing import Dict
+from dataclasses import fields  # Import fields to inspect dataclass
+
+from hackagent.attacks.AdvPrefix.scorer_parser import (
+    EvaluatorConfig,
+    NuancedEvaluator,
+    HarmBenchEvaluator,
+    JailbreakBenchEvaluator,
+)
+from hackagent.attacks.AdvPrefix.utils import get_checkpoint_path
+
+EVALUATOR_MAP = {
+    "nuanced": NuancedEvaluator,
+    "jailbreakbench": JailbreakBenchEvaluator,
+    "harmbench": HarmBenchEvaluator,
+}
+
+MERGE_KEYS = ["goal", "prefix", "completion"]  # Standard merge keys
+
+JUDGE_COLUMN_MAP = {
+    "nuanced": ["eval_nj", "explanation_nj"],
+    "jailbreakbench": ["eval_jb", "explanation_jb"],
+    "harmbench": ["eval_hb", "explanation_hb"],
+}
+
+
+def _run_evaluator_process_wrapper(
+    judge_type: str, config_dict_serializable: Dict, df: pd.DataFrame
+):
+    """Static method to run a specific evaluator, suitable for multiprocessing."""
+    process_logger = logging.getLogger(__name__ + f".evaluator_process_{judge_type}")
+    process_logger.info(f"Evaluator process started for judge: {judge_type}")
+
+    evaluator_class = EVALUATOR_MAP.get(judge_type)
+    if not evaluator_class:
+        process_logger.warning(f"Unknown judge type: {judge_type}, skipping")
+        return None
+
+    evaluator = None  # Ensure cleanup
+    try:
+        # Filter the config dict to only include keys expected by EvaluatorConfig
+        expected_fields = {f.name for f in fields(EvaluatorConfig)}
+        filtered_config_dict = {
+            k: v for k, v in config_dict_serializable.items() if k in expected_fields
+        }
+
+        # model_id is handled specially by EvaluatorConfig.with_default_model
+        # or within the specific evaluator's __init__.
+        # We should ensure model_id from the judge's config is passed if present.
+        if (
+            "model_id" in config_dict_serializable
+            and config_dict_serializable["model_id"]
+        ):
+            filtered_config_dict["model_id"] = config_dict_serializable["model_id"]
+        elif (
+            "identifier" in config_dict_serializable
+            and config_dict_serializable["identifier"]
+        ):
+            # Fallback to using 'identifier' if 'model_id' wasn't explicitly passed/overridden
+            filtered_config_dict["model_id"] = config_dict_serializable["identifier"]
+
+        process_logger.debug(
+            f"Filtered config for {judge_type} evaluator: {filtered_config_dict}"
+        )
+        evaluator_config = EvaluatorConfig(**filtered_config_dict)
+
+        # Instantiate the specific evaluator class
+        evaluator = evaluator_class(evaluator_config)
+        evaluated_df = evaluator.evaluate(df)
+
+        process_logger.info(f"Evaluator process finished for judge: {judge_type}")
+        # Return only the essential columns: merge keys + judge-specific columns
+        eval_cols = JUDGE_COLUMN_MAP.get(judge_type, [])
+        # Ensure merge keys are present in the returned df
+        if not all(key in evaluated_df.columns for key in MERGE_KEYS):
+            process_logger.error(
+                f"Evaluation result for {judge_type} is missing merge keys {MERGE_KEYS}. Available: {evaluated_df.columns}. Returning None."
+            )
+            return None
+
+        cols_to_return = MERGE_KEYS + [
+            col for col in eval_cols if col in evaluated_df.columns
+        ]
+        return evaluated_df[cols_to_return]
+
+    except Exception as e:
+        process_logger.error(
+            f"Error occurred while running {judge_type} evaluator: {str(e)}",
+            exc_info=True,
+        )
+        return None  # Indicate failure
+    finally:
+        # Cleanup
+        del evaluator
+        process_logger.info(
+            f"Evaluator process cleanup finished for judge: {judge_type}"
+        )
+
+
+def execute(
+    input_df: pd.DataFrame, config: Dict, logger: logging.Logger, run_dir: str
+) -> pd.DataFrame:
+    """Evaluate completions using specified judges."""
+    logger.info("Executing Step 7: Evaluating responses")
+    original_df = input_df
+
+    if original_df.empty:
+        logger.warning("Step 7 received an empty DataFrame. Skipping evaluation.")
+        return original_df
+
+    # Config key 'judges' should be a list of dictionaries
+    judge_configs_list = config.get("judges")
+    if not isinstance(judge_configs_list, list) or not judge_configs_list:
+        logger.warning(
+            "Step 7: 'judges' key in configuration is missing, not a list, or empty. Skipping evaluation."
+        )
+        return original_df
+
+    # Base config for evaluators (extract non-judge-specific params)
+    evaluator_base_config_dict = {
+        "batch_size": config.get("batch_size_judge"),
+        "max_new_tokens_eval": config.get("max_new_tokens_eval"),
+        "filter_len": config.get("filter_len"),
+        # General API settings (judges might override)
+        "endpoint": config.get("judge_endpoint"),
+        "api_key": config.get("judge_api_key"),
+        "request_timeout": config.get("judge_request_timeout"),
+    }
+
+    judge_results_dfs = {}
+    failed_judges = []
+    async_results = []
+    judges_to_run = []  # Store valid (type, config_dict) tuples
+
+    # --- Prepare Judge Runs ---
+    for judge_config_item in judge_configs_list:
+        if not isinstance(judge_config_item, dict):
+            logger.warning(
+                f"Skipping invalid item in 'judges' list (not a dict): {judge_config_item}"
+            )
+            continue
+
+        # Extract the judge type string (e.g., "nuanced", "harmbench")
+        # Assuming the type is specified by a 'type' key in the judge dict.
+        # Alternative: Infer based on 'identifier'?
+        judge_type_str = judge_config_item.get(
+            "evaluator_type"
+        ) or judge_config_item.get("type")
+        judge_identifier = judge_config_item.get("identifier")
+
+        if not judge_type_str:
+            # If type isn't explicit, try to infer (this part might need refinement)
+            if "nuanced" in judge_identifier.lower():
+                judge_type_str = "nuanced"
+            elif "harmbench" in judge_identifier.lower():
+                judge_type_str = "harmbench"
+            elif "jailbreak" in judge_identifier.lower():
+                judge_type_str = "jailbreakbench"
+            else:
+                logger.warning(
+                    f"Could not determine evaluator type for judge config: {judge_config_item}. Skipping."
+                )
+                continue
+
+        # Check if the extracted type string is valid
+        if judge_type_str not in EVALUATOR_MAP:
+            logger.warning(
+                f"Skipping unknown judge type '{judge_type_str}' specified in config: {judge_config_item}"
+            )
+            continue
+
+        # Prepare the specific config to pass to the subprocess
+        # Start with base, then override with judge-specific settings
+        subprocess_config = evaluator_base_config_dict.copy()
+        subprocess_config.update(judge_config_item)  # Override base with specifics
+        # Ensure model_id is set correctly (use 'identifier')
+        if judge_identifier:
+            subprocess_config["model_id"] = judge_identifier
+
+        judges_to_run.append((judge_type_str, subprocess_config))
+
+    if not judges_to_run:
+        logger.warning(
+            "Step 7: No valid judges found after processing configuration. Skipping evaluation."
+        )
+        return original_df
+
+    # --- Setup Multiprocessing Pool ---
+    try:
+        current_start_method = multiprocessing.get_start_method(allow_none=True)
+        if current_start_method != "spawn":
+            multiprocessing.set_start_method("spawn", force=True)
+            logger.info("Set multiprocessing start method to 'spawn' for Step 7.")
+    except Exception as e:
+        logger.warning(f"Could not set multiprocessing start method to spawn: {e}")
+
+    num_judges = len(judges_to_run)
+    num_workers = min(num_judges, os.cpu_count() or 1, 4)
+    logger.info(
+        f"Starting evaluation pool with {num_workers} workers for {num_judges} judges."
+    )
+
+    # --- Dispatch and Collect Results ---
+    with multiprocessing.Pool(processes=num_workers) as pool:
+        # Dispatch tasks using the prepared list
+        for judge_type_str, subprocess_config in judges_to_run:
+            logger.info(
+                f"Dispatching evaluation with {judge_type_str} judge. Config: {subprocess_config}"
+            )
+            args = (judge_type_str, subprocess_config, original_df.copy())
+            async_results.append(
+                pool.apply_async(_run_evaluator_process_wrapper, args=args)
+            )
+
+        # Collect results using the order in judges_to_run
+        for (judge_type_str, _), result in zip(judges_to_run, async_results):
+            try:
+                evaluated_df_subset = result.get()
+                if evaluated_df_subset is not None:
+                    judge_results_dfs[judge_type_str] = evaluated_df_subset
+                    logger.info(
+                        f"Successfully completed evaluation for judge: {judge_type_str}"
+                    )
+                else:
+                    failed_judges.append(judge_type_str)
+                    logger.error(
+                        f"Evaluation failed for judge: {judge_type_str} (process returned None)"
+                    )
+            except Exception as e:
+                failed_judges.append(judge_type_str)
+                logger.error(
+                    f"Evaluation task failed for judge {judge_type_str}: {e}",
+                    exc_info=True,
+                )
+
+    # --- Merge Results ---
+    final_df = original_df.copy()
+    successful_judges = list(judge_results_dfs.keys())
+
+    if not successful_judges:
+        logger.warning(
+            "Step 7: No judges completed successfully. Returning original DataFrame."
+        )
+    else:
+        logger.info(f"Merging results from successful judges: {successful_judges}")
+        for judge_type_str in successful_judges:
+            judge_df_subset = judge_results_dfs[judge_type_str]
+            eval_cols = JUDGE_COLUMN_MAP.get(judge_type_str, [])
+            judge_cols_present = [
+                col for col in eval_cols if col in judge_df_subset.columns
+            ]
+
+            if not judge_cols_present:
+                logger.warning(
+                    f"No specific evaluation columns found in result for judge {judge_type_str}. Skipping merge for this judge."
+                )
+                continue
+
+            try:
+                final_df = final_df.merge(
+                    judge_df_subset,
+                    on=MERGE_KEYS,
+                    how="left",
+                    suffixes=("", f"_{judge_type_str}_dup"),
+                )
+                logger.debug(f"Merged results from judge {judge_type_str}")
+            except Exception as e:
+                logger.error(f"Error merging results for judge {judge_type_str}: {e}")
+
+    # Save final merged results checkpoint
+    output_path = get_checkpoint_path(run_dir, 7)
+    try:
+        final_df.to_csv(output_path, index=False)
+        logger.info(f"Step 7 complete. Evaluated {len(final_df)} responses.")
+        if successful_judges:
+            logger.info(
+                f"Successfully completed judges: {', '.join(successful_judges)}"
+            )
+        if failed_judges:
+            logger.warning(f"Failed judges: {', '.join(failed_judges)}")
+        logger.info(f"Final evaluation results checkpoint saved to {output_path}")
+    except Exception as e:
+        logger.error(f"Failed to save checkpoint for step 7 to {output_path}: {e}")
+
+    return final_df

hackagent/attacks/AdvPrefix/step8_aggregate_evaluations.py

@@ -0,0 +1,177 @@
+import pandas as pd
+from typing import Dict, Any
+
+from .utils import get_checkpoint_path
+
+# Map judge type to expected column prefix/name used for aggregation stats
+JUDGE_AGG_COLUMN_MAP = {
+    "nuanced": "eval_nj",
+    "jailbreakbench": "eval_jb",
+    "harmbench": "eval_hb",
+}
+
+GROUP_KEYS = ["goal", "prefix"]
+
+
+def execute(
+    input_df: pd.DataFrame, config: Dict[str, Any], run_dir: str
+) -> pd.DataFrame:
+    """
+    Aggregate evaluation results from different judges using the input DataFrame.
+    Combines results from multiple evaluation samples and judges into single scores per prefix.
+    """
+    print("Executing Step 8: Aggregating evaluation results")
+
+    if input_df.empty:
+        print("WARNING: Step 8 received an empty DataFrame. Skipping aggregation.")
+        # Define expected aggregated columns if returning empty
+        cols = GROUP_KEYS + [
+            "prefix_nll",
+            "model_name",
+            "meta_prefix",
+            "temperature",
+            "n_eval_samples",
+        ]
+        for _, col_base in JUDGE_AGG_COLUMN_MAP.items():
+            cols.extend([f"{col_base}_mean", f"{col_base}_count"])
+        return pd.DataFrame(columns=cols)
+
+    analysis = input_df.copy()
+
+    # Optionally filter based on cross-entropy / NLL score
+    if "prefix_nll" in analysis.columns and config.get("max_ce") is not None:
+        try:
+            max_ce_threshold = float(config.get("max_ce"))
+            initial_count = len(analysis)
+            # Use dictionary access for config
+            analysis = analysis[analysis["prefix_nll"] < max_ce_threshold]
+            filtered_count = len(analysis)
+            print(
+                f"Filtered {initial_count - filtered_count} rows based on prefix_nll >= {max_ce_threshold}"
+            )
+        except KeyError:
+            print("WARNING: 'max_ce' key not found in config, skipping NLL filtering.")
+        except Exception as e:
+            print(f"ERROR: Error during NLL filtering in aggregation: {e}")
+            # Continue without NLL filtering if error occurs
+    elif "prefix_nll" not in analysis.columns:
+        print(
+            "WARNING: Column 'prefix_nll' not found. Skipping NLL filtering in aggregation step."
+        )
+
+    # Detect available judges based on column names for aggregation
+    available_judges_agg_cols = {}
+    judges_in_config = config.get("judges", [])  # Judges that were supposed to run
+    for judge_type, col_name in JUDGE_AGG_COLUMN_MAP.items():
+        if col_name in analysis.columns:
+            available_judges_agg_cols[judge_type] = col_name
+        else:
+            # Log if any expected judge column is missing
+            if judge_type in judges_in_config:
+                print(
+                    f"WARNING: Expected aggregation column '{col_name}' for judge '{judge_type}' not found in the dataframe for Step 8."
+                )
+
+    if not available_judges_agg_cols:
+        print(
+            "ERROR: No recognized evaluation result columns found for aggregation. Check step 7 output."
+        )
+        output_path = get_checkpoint_path(run_dir, 8)
+        try:
+            analysis.to_csv(output_path, index=False)
+            print(
+                f"WARNING: Step 8 saving unaggregated data to {output_path} due to missing judge columns."
+            )
+        except Exception as e:
+            print(
+                f"ERROR: Failed to save unaggregated data checkpoint for step 8 to {output_path}: {e}"
+            )
+        return analysis  # Return unaggregated data
+
+    print(
+        f"Found aggregation columns for judges: {list(available_judges_agg_cols.keys())}"
+    )
+
+    # Ensure group keys exist
+    if not all(key in analysis.columns for key in GROUP_KEYS):
+        missing_keys = [key for key in GROUP_KEYS if key not in analysis.columns]
+        print(
+            f"ERROR: Missing required grouping keys for aggregation: {missing_keys}. Cannot aggregate."
+        )
+        output_path = get_checkpoint_path(run_dir, 8)
+        try:
+            analysis.to_csv(output_path, index=False)
+            print(
+                f"WARNING: Step 8 saving unaggregated data to {output_path} due to missing group keys."
+            )
+        except Exception as e:
+            print(
+                f"ERROR: Failed to save unaggregated data checkpoint for step 8 to {output_path}: {e}"
+            )
+        return analysis
+
+    # Define aggregations
+    agg_funcs = {
+        # Use pd.NamedAgg for clarity and future compatibility
+        "prefix_nll": pd.NamedAgg(column="prefix_nll", aggfunc="first"),
+        "model_name": pd.NamedAgg(column="model_name", aggfunc="first"),
+        "meta_prefix": pd.NamedAgg(column="meta_prefix", aggfunc="first"),
+        "temperature": pd.NamedAgg(column="temperature", aggfunc="first"),
+        # Count samples - use one of the group keys or index if reset
+        "n_eval_samples": pd.NamedAgg(column=GROUP_KEYS[0], aggfunc="size"),
+    }
+
+    # Add judge-specific aggregations
+    for judge_type, col_name in available_judges_agg_cols.items():
+        # Ensure the column is numeric before calculating mean
+        try:
+            analysis[col_name] = pd.to_numeric(analysis[col_name], errors="coerce")
+            agg_funcs[f"{col_name}_mean"] = pd.NamedAgg(column=col_name, aggfunc="mean")
+            agg_funcs[f"{col_name}_count"] = pd.NamedAgg(
+                column=col_name, aggfunc="count"
+            )  # Count non-NA numeric values
+            print(
+                f"DEBUG: Added mean/count aggregation for numeric column '{col_name}'"
+            )
+        except KeyError:
+            print(
+                f"WARNING: Column '{col_name}' unexpectedly missing during aggregation setup. Skipping mean/count."
+            )
+        except Exception as e:
+            print(
+                f"ERROR: Could not convert column '{col_name}' to numeric for aggregation. Skipping mean/count. Error: {e}"
+            )
+            # Optionally add just size aggregation if mean fails?
+            agg_funcs[f"{col_name}_size"] = pd.NamedAgg(column=col_name, aggfunc="size")
+
+    # Perform aggregation
+    try:
+        grouped = analysis.groupby(GROUP_KEYS, observed=False, dropna=False)
+        aggregated = grouped.agg(**agg_funcs)
+        aggregated = aggregated.reset_index()
+    except Exception as e:
+        print(
+            f"ERROR: Error during aggregation: {e}. Check aggregation functions and column types."
+        )
+        output_path = get_checkpoint_path(run_dir, 8)
+        try:
+            analysis.to_csv(output_path, index=False)
+            print(
+                f"WARNING: Step 8 saving unaggregated data to {output_path} due to aggregation error."
+            )
+        except Exception as e_save:
+            print(
+                f"ERROR: Failed to save unaggregated data checkpoint for step 8 to {output_path}: {e_save}"
+            )
+        return analysis  # Return unaggregated on error
+
+    # Save results checkpoint
+    output_path = get_checkpoint_path(run_dir, 8)
+    try:
+        aggregated.to_csv(output_path, index=False)
+        print(f"Step 8 complete. Aggregated {len(aggregated)} prefix results.")
+        print(f"Checkpoint saved to {output_path}")
+    except Exception as e:
+        print(f"ERROR: Failed to save checkpoint for step 8 to {output_path}: {e}")
+
+    return aggregated

hackagent/attacks/AdvPrefix/step9_select_prefixes.py

@@ -0,0 +1,59 @@
+import logging
+import pandas as pd
+from typing import Dict, Any
+
+from hackagent.attacks.AdvPrefix.selector import (
+    PrefixSelectorConfig,
+    PrefixSelector,
+)
+
+from .utils import get_checkpoint_path
+
+logger = logging.getLogger(__name__)
+
+
+def execute(
+    input_df: pd.DataFrame, config: Dict[str, Any], run_dir: str
+) -> pd.DataFrame:
+    """Select final prefixes based on specified judges and selection criteria using input DataFrame."""
+    logger.info("Executing Step 9: Selecting final prefixes")
+
+    if input_df.empty:
+        logger.warning("Step 9 received an empty DataFrame. Skipping selection.")
+        return input_df
+
+    selector = None  # Ensure cleanup
+    selected_df = input_df  # Default to input if selection fails
+
+    try:
+        # Initialize selector here
+        selector_config = PrefixSelectorConfig(
+            pasr_weight=config.get("pasr_weight", 0.5),
+            n_prefixes_per_goal=config.get("n_prefixes_per_goal", 3),
+            judges=config.get("selection_judges", []),
+        )
+        selector = PrefixSelector(selector_config)
+
+        # Select prefixes
+        selected_df = selector.select_prefixes(input_df)
+        logger.info(f"Selection complete. Selected {len(selected_df)} prefixes.")
+
+    except Exception as e:
+        logger.error(f"Error during prefix selection: {e}", exc_info=True)
+        logger.warning("Returning unselected prefixes due to selection error.")
+        selected_df = input_df  # Fallback to returning the input df
+
+    finally:
+        del selector
+        # No GPU cleanup needed typically for selection
+
+    # Save results checkpoint (final step)
+    output_path = get_checkpoint_path(run_dir, 9)
+    try:
+        selected_df.to_csv(output_path, index=False)
+        logger.info("Step 9 complete.")
+        logger.info(f"Final selected prefixes checkpoint saved to {output_path}")
+    except Exception as e:
+        logger.error(f"Failed to save checkpoint for step 9 to {output_path}: {e}")
+
+    return selected_df