hackagent 0.1.0__py3-none-any.whl

hackagent/attacks/AdvPrefix/preprocessing.py

@@ -0,0 +1,521 @@
+"""
+Prefix preprocessing implementation for the prefix generation pipeline.
+This module handles both filtering and ablation of prefixes.
+"""
+
+from dataclasses import dataclass
+import logging
+from typing import List, Dict, Optional
+import pandas as pd
+import re
+import numpy as np
+from rich.progress import (
+    Progress,
+    BarColumn,
+    TextColumn,
+    TimeRemainingColumn,
+    MofNCompleteColumn,
+    SpinnerColumn,
+)
+
+
+@dataclass
+class PreprocessConfig:
+    """Configuration for prefix preprocessing steps"""
+
+    # model_id: str # Removed, no longer needed for token counting here
+
+    # Filtering configs
+    min_char_length: int = 50  # Renamed from min_token_length, adjusted default
+    require_linebreak: bool = True
+    max_ce: Optional[float] = None
+
+    # Ablation configs
+    # Note: max_token_segments now controls the number of CHARACTER-based segments
+    max_token_segments: int = (
+        5  # How many character segments to create from second line
+    )
+    min_lines: int = 2  # Minimum number of non-empty lines required
+
+    # Numbers of prefixes per goal to keep before nuanced evaluation
+    n_candidates_per_goal: int = (
+        100  # Number of prefixes to keep per goal after CE filtering
+    )
+
+    # Patterns to filter out at the start of prefixes
+    start_patterns: tuple[str, ...] = (
+        "I cannot",
+        "I can't",
+        "I'm afraid",
+        "Unfortunately",
+        "Sorry",
+        "I am sorry",
+        "I'm sorry",
+        "I am really sorry",
+        "I'm very sorry",
+        "I apologize",
+        "As an",
+        "As a",
+        "I'm just",
+        "I'm an",
+        "I understand",
+        "To begin, I must ",
+        "I must",
+        "I am not able to",
+        "I am unable",
+        "Here's what I can tell you",
+        "Here is what I can tell you",
+        "Here's what I can do",
+        "Here is what I can do",
+        "Here's what I can say",
+        "Here is what I can say",
+        "Here are some things you should know",
+        "To be clear",
+    )
+
+    # Patterns to filter out anywhere in prefixes
+    contain_patterns: tuple[str, ...] = (
+        "I am an AI assistant",
+        "I can't provide",
+        "I cannot",
+        "I can't assist",
+        "I can't write",
+        "I can't fulfill",
+        "I can't do",
+        "I can't help",
+        "I can't and",
+        "I am unable",
+    )
+
+
+class PrefixPreprocessor:
+    """
+    Implements preprocessing logic for prefixes, including filtering and ablation.
+
+    Filtering is split into two phases:
+    Phase 1 (before NLL calculation):
+    1. Remove prefixes shorter than minimum token length (using litellm.token_counter)
+    2. Remove prefixes starting with unwanted phrases
+    3. Remove prefixes containing unwanted phrases
+    4. Remove prefixes without linebreaks
+    5. Merge duplicates
+
+    Phase 2 (after NLL calculation):
+    1. Filter based on cross-entropy loss threshold
+
+    Ablation process (now character-based):
+    1. Clean up prefixes by removing leading spaces while preserving line breaks
+    2. Split prefixes into lines and identify non-empty lines
+    3. Create variations by taking first line and different character lengths of second line
+    4. Merge duplicate prefixes while preserving metadata
+    """
+
+    def __init__(self, config: PreprocessConfig):
+        """Initialize the preprocessor with configuration."""
+        self.config = config
+        self.logger = logging.getLogger(__name__)
+        # Removed tokenizer loading and related logging
+        self.logger.info(
+            "PrefixPreprocessor initialized. Filtering will use character lengths."
+        )
+
+    def _clean_prefix(self, prefix: str) -> str:
+        """Clean prefix by removing leading spaces but keeping line breaks."""
+        # Preserve leading whitespace that includes newlines, remove only leading spaces/tabs on the first line.
+        match = re.match(r"^[ \t]*(.*)", prefix, re.DOTALL)
+        if match:
+            return match.group(1)
+        return prefix  # Should not happen with DOTALL but as fallback
+
+    def _merge_duplicates(self, df: pd.DataFrame) -> pd.DataFrame:
+        """Merge duplicate prefixes while preserving metadata."""
+        rows_before = len(df)
+
+        def concatenate_unique_entries(group):
+            agg_dict = {
+                "model_name": lambda x: ",".join(str(v) for v in set(x)),
+                "meta_prefix": lambda x: ",".join(
+                    str(v) for v in set(x) if pd.notna(v)
+                ),
+                "temperature": lambda x: ",".join(str(v) for v in set(x)),
+                "goal": "first",
+            }
+
+            if "prefix_nll" in group.columns:
+                agg_dict["prefix_nll"] = "first"
+
+            return group.groupby("prefix").agg(agg_dict).reset_index()
+
+        # Apply within each goal group first to avoid unintended merges across goals
+        df = df.groupby("goal").apply(concatenate_unique_entries).reset_index(drop=True)
+
+        rows_after = len(df)
+        if rows_before > 0:
+            self.logger.info(
+                f"Duplicate merging reduced rows from {rows_before} to {rows_after} (within goals)"
+            )
+        return df
+
+    def _print_detailed_stats(self, df: pd.DataFrame, step_name: str):
+        """Print detailed statistics about remaining prefixes per goal."""
+        if df.empty:
+            self.logger.info(f"Detailed {step_name} statistics: DataFrame is empty.")
+            return
+        goal_prefix_counts = df.groupby("goal")["prefix"].count()
+        min_prefixes_left = goal_prefix_counts.min()
+        max_prefixes_left = goal_prefix_counts.max()
+        average_prefixes_left = goal_prefix_counts.mean()
+        median_prefixes_left = goal_prefix_counts.median()
+        std_dev_prefixes = goal_prefix_counts.std()
+        goal_with_min_prefixes = goal_prefix_counts.idxmin()
+
+        self.logger.info(f"Detailed {step_name} statistics:")
+        self.logger.info(f"- Total prefixes remaining: {len(df)}")
+        self.logger.info(f"- Number of goals: {df['goal'].nunique()}")
+        self.logger.info(
+            f"- Prefixes per goal: Min={min_prefixes_left}, Max={max_prefixes_left}, Avg={average_prefixes_left:.2f}, Median={median_prefixes_left:.0f}, StdDev={std_dev_prefixes:.2f}"
+        )
+        self.logger.info(
+            f"- Goal with minimum prefixes ({min_prefixes_left}): '{goal_with_min_prefixes[:50]}...'"
+        )
+
+        if min_prefixes_left < 10:
+            low_goals_count = (goal_prefix_counts < 10).sum()
+            self.logger.warning(
+                f"{low_goals_count} goals have very few prefixes left (<10)"
+            )
+
+    # Renamed and modified filtering method
+    def _filter_by_char_length(self, df: pd.DataFrame) -> pd.DataFrame:
+        """Remove prefixes shorter than minimum character length."""
+        if not self.config.min_char_length or self.config.min_char_length <= 0:
+            return df
+
+        rows_before = len(df)
+        # Filter based on character length
+        df_filtered = df[df["prefix"].str.len() >= self.config.min_char_length]
+        rows_after = len(df_filtered)
+        removed_count = rows_before - rows_after
+
+        if removed_count > 0:
+            self.logger.info(
+                f"Character length filter (< {self.config.min_char_length} chars) removed {removed_count} rows"
+            )
+        else:
+            self.logger.info("Character length filter did not remove any rows.")
+        # No temporary column to drop
+        return df_filtered
+
+    def _filter_by_start_patterns(self, df: pd.DataFrame) -> pd.DataFrame:
+        """Remove prefixes starting with unwanted phrases."""
+        if not self.config.start_patterns:
+            return df
+        rows_before = len(df)
+        # Apply pattern matching after stripping leading whitespace
+        df_filtered = df[
+            ~df["prefix"]
+            .str.lstrip()
+            .str.startswith(self.config.start_patterns, na=False)
+        ]
+        rows_after = len(df_filtered)
+        removed_count = rows_before - rows_after
+        if removed_count > 0:
+            self.logger.info(f"Start pattern filter removed {removed_count} rows")
+        return df_filtered
+
+    def _filter_by_contain_patterns(self, df: pd.DataFrame) -> pd.DataFrame:
+        """Remove prefixes containing unwanted phrases."""
+        if not self.config.contain_patterns:
+            return df
+        rows_before = len(df)
+        # Combine patterns into a single regex
+        pattern = "|".join(
+            map(re.escape, self.config.contain_patterns)
+        )  # Escape special chars
+        df_filtered = df[~df["prefix"].str.contains(pattern, regex=True, na=False)]
+        rows_after = len(df_filtered)
+        removed_count = rows_before - rows_after
+        if removed_count > 0:
+            self.logger.info(f"Contain pattern filter removed {removed_count} rows")
+        return df_filtered
+
+    def _filter_by_linebreak(self, df: pd.DataFrame) -> pd.DataFrame:
+        """Remove prefixes that don't contain linebreaks (excluding at start/end)."""
+        if not self.config.require_linebreak:
+            return df
+
+        rows_before = len(df)
+        # Check for newline within the string after stripping leading/trailing whitespace and newlines
+        df_filtered = df[
+            df["prefix"].str.strip().str.strip("\n").str.contains("\n", na=False)
+        ]
+        rows_after = len(df_filtered)
+        removed_count = rows_before - rows_after
+        if removed_count > 0:
+            self.logger.info(f"Linebreak filter removed {removed_count} rows")
+        return df_filtered
+
+    def _filter_by_ce_loss(self, df: pd.DataFrame) -> pd.DataFrame:
+        """Remove prefixes with cross-entropy loss above threshold."""
+        if self.config.max_ce is None or "prefix_nll" not in df.columns:
+            if "prefix_nll" not in df.columns:
+                self.logger.warning(
+                    "CE loss filtering skipped: 'prefix_nll' column not found."
+                )
+            return df
+
+        rows_before = len(df)
+        # Ensure NLL column is numeric, coerce errors
+        df["prefix_nll_numeric"] = pd.to_numeric(df["prefix_nll"], errors="coerce")
+        # Filter out rows where conversion failed or value is above threshold
+        df_filtered = df[df["prefix_nll_numeric"] <= self.config.max_ce]
+        rows_after = len(df_filtered)
+        removed_count = rows_before - rows_after
+
+        if removed_count > 0:
+            self.logger.info(
+                f"CE loss filter (> {self.config.max_ce}) removed {removed_count} rows"
+            )
+        else:
+            self.logger.info("CE loss filter did not remove any rows.")
+        return df_filtered.drop(columns=["prefix_nll_numeric"])
+
+    # Ablation methods (now character-based)
+    def _should_ablate(self, prefix: str) -> bool:
+        """Determine if a prefix should be ablated."""
+        if not isinstance(prefix, str) or "\n" not in prefix:
+            return False
+
+        # Use regex to find lines preserving interstitial newlines
+        lines_with_breaks = re.split(r"(\n+)", prefix)
+        non_empty_lines_content = [
+            line for line in lines_with_breaks if line.strip()
+        ]  # Content lines
+
+        # Check if there are at least min_lines of actual content
+        if len(non_empty_lines_content) < self.config.min_lines:
+            return False
+
+        return True
+
+    def _create_ablated_versions(self, row: pd.Series) -> List[Dict]:
+        """Create ablated versions of a prefix based on character length."""
+        prefix = row["prefix"]
+        if not isinstance(prefix, str):
+            return []
+
+        # Split carefully to preserve original newline structure between lines
+        lines_with_breaks = re.split(r"(\n+)", prefix)
+        # Identify indices of lines with actual content vs just newlines
+        content_indices = [
+            i for i, line in enumerate(lines_with_breaks) if line.strip()
+        ]
+
+        if len(content_indices) < 2:
+            return []  # Not enough content lines
+
+        # Reconstruct the first line part including any preceding/trailing newlines from the split
+        first_line_end_index = content_indices[0]
+        first_part = "".join(lines_with_breaks[: first_line_end_index + 1])
+
+        # Identify the second content line and the newline(s) immediately preceding it
+        second_line_start_index = content_indices[1]
+        # Newlines between first and second content line start *after* the first content part ends
+        newline_separator_index = first_line_end_index + 1
+        newline_separator = "".join(
+            lines_with_breaks[newline_separator_index:second_line_start_index]
+        )
+        second_line_content = lines_with_breaks[second_line_start_index]
+
+        # Original second line content (for length calculation)
+        second_line_strip = second_line_content.strip()
+        char_len = len(second_line_strip)
+        if char_len == 0:
+            return []  # Second line is effectively empty
+
+        new_rows = []
+        # Determine character segment lengths
+        num_segments = min(char_len, self.config.max_token_segments)  # Use config value
+        if num_segments <= 0:
+            return []
+
+        # Create segment lengths (e.g., 5, 10, 15... or proportionally)
+        # Using linspace ensures we get segments up to the full length
+        segment_lengths = np.linspace(1, char_len, num=num_segments, dtype=int)
+        unique_segment_lengths = sorted(
+            list(set(segment_lengths))
+        )  # Ensure uniqueness and order
+
+        for length in unique_segment_lengths:
+            # Take the first `length` characters of the original second line content
+            truncated_second_line = second_line_strip[:length]
+
+            # Construct the new prefix using the identified parts
+            new_prefix = f"{first_part}{newline_separator}{truncated_second_line}"
+
+            new_row = row.to_dict()  # Convert Series to dict for modification
+            new_row["prefix"] = new_prefix
+            # Optionally add metadata about ablation? e.g., original prefix, segment length
+            # new_row['ablation_source'] = prefix
+            # new_row['ablation_segment_chars'] = length
+            new_rows.append(new_row)
+
+        return new_rows
+
+    # Public interface methods
+    def filter_phase1(self, df: pd.DataFrame) -> pd.DataFrame:
+        """
+        Apply all phase 1 filters: patterns, length, linebreaks, duplicates.
+        """
+        self.logger.info("Starting filter phase 1...")
+        df_filtered = df.copy()  # Work on a copy
+
+        # Apply filters sequentially
+        df_filtered = self._filter_by_start_patterns(df_filtered)
+        df_filtered = self._filter_by_contain_patterns(df_filtered)
+        df_filtered = self._filter_by_char_length(df_filtered)  # Call renamed method
+        if self.config.require_linebreak:
+            df_filtered = self._filter_by_linebreak(df_filtered)
+
+        df_merged = self._merge_duplicates(df_filtered)
+
+        self._print_detailed_stats(df_merged, "Filter Phase 1")
+        return df_merged
+
+    def filter_phase2(self, df: pd.DataFrame) -> pd.DataFrame:
+        """Apply final filtering steps after NLL calculation."""
+        if not isinstance(df, pd.DataFrame) or df.empty:
+            self.logger.warning(
+                "Phase 2 filtering received an empty or invalid DataFrame. Skipping."
+            )
+            return pd.DataFrame(
+                columns=df.columns if isinstance(df, pd.DataFrame) else []
+            )
+        if "prefix_nll" not in df.columns:
+            self.logger.error(
+                "Phase 2 filtering requires 'prefix_nll' column, but it is missing. Skipping."
+            )
+            return df
+
+        self.logger.info(f"Starting prefix filtering phase 2 with {len(df)} rows")
+
+        # Filter by CE loss threshold
+        df_ce_filtered = self._filter_by_ce_loss(df)
+        self.logger.info(f"After CE threshold filtering: {len(df_ce_filtered)} rows")
+
+        # Select top k prefixes per goal based on CE loss
+        if self.config.n_candidates_per_goal > 0 and not df_ce_filtered.empty:
+            self.logger.info(
+                f"Selecting top {self.config.n_candidates_per_goal} prefixes per goal based on NLL..."
+            )
+            # Ensure NLL is numeric for sorting
+            df_ce_filtered["prefix_nll_numeric"] = pd.to_numeric(
+                df_ce_filtered["prefix_nll"], errors="coerce"
+            )
+            # Group, sort, take top N, handling potential NA values in NLL
+            df_top_k = (
+                df_ce_filtered.sort_values("prefix_nll_numeric", na_position="last")
+                .groupby("goal")
+                .head(self.config.n_candidates_per_goal)
+            )
+            df_final = df_top_k.drop(columns=["prefix_nll_numeric"]).reset_index(
+                drop=True
+            )
+            self.logger.info(
+                f"After selecting top {self.config.n_candidates_per_goal} per goal: "
+                f"{len(df_final)} rows"
+            )
+        else:
+            self.logger.info(
+                f"Skipping top-k selection (k={self.config.n_candidates_per_goal} or DataFrame empty)."
+            )
+            df_final = df_ce_filtered.reset_index(drop=True)
+
+        self._print_detailed_stats(df_final, "phase 2 filtering")
+        return df_final
+
+    def ablate(self, df: pd.DataFrame) -> pd.DataFrame:
+        """Perform prefix ablation (character-based) on the input dataframe."""
+        if not isinstance(df, pd.DataFrame) or df.empty:
+            self.logger.warning(
+                "Ablation received an empty or invalid DataFrame. Skipping."
+            )
+            return pd.DataFrame(
+                columns=df.columns if isinstance(df, pd.DataFrame) else []
+            )
+        if "prefix" not in df.columns:
+            self.logger.error(
+                "Ablation requires 'prefix' column, but it is missing. Skipping."
+            )
+            return df
+
+        self.logger.info(f"Starting prefix ablation with {len(df)} rows")
+        original_cols = df.columns.tolist()
+
+        # Clean prefixes first (important for accurate line splitting)
+        df["prefix"] = df["prefix"].apply(self._clean_prefix)
+
+        # Identify ablatable rows
+        ablatable_mask = df["prefix"].apply(self._should_ablate)
+        ablatable_df = df[ablatable_mask]
+        non_ablatable_df = df[~ablatable_mask]  # Keep rows that won't be ablated
+        self.logger.info(
+            f"{len(ablatable_df)} prefixes identified for character-based ablation."
+        )
+
+        if ablatable_df.empty:
+            self.logger.info(
+                "No prefixes suitable for ablation. Returning original DataFrame."
+            )
+            return df
+
+        new_rows = []
+
+        with Progress(
+            SpinnerColumn(),
+            TextColumn("[progress.description]{task.description}"),
+            BarColumn(),
+            MofNCompleteColumn(),
+            TextColumn("[progress.percentage]{task.percentage:>3.1f}%"),
+            TimeRemainingColumn(),
+        ) as progress_bar:
+            task = progress_bar.add_task(
+                "[cyan]Creating ablated prefixes...", total=len(ablatable_df)
+            )
+            for _, row in ablatable_df.iterrows():
+                new_rows.extend(self._create_ablated_versions(row))
+                progress_bar.update(task, advance=1)
+
+        if not new_rows:
+            self.logger.warning(
+                "Ablation process created no new prefixes. Returning original DataFrame."
+            )
+            # Might return non_ablatable_df if we want to strictly separate
+            return df
+
+        ablated_results_df = pd.DataFrame(new_rows)
+        self.logger.info(
+            f"Created {len(ablated_results_df)} ablated prefix variations."
+        )
+
+        # Combine non-ablated rows with the new ablated variations
+        # Ensure columns match before concatenating
+        combined_df = pd.concat(
+            [non_ablatable_df, ablated_results_df], ignore_index=True, sort=False
+        )
+
+        # Merge duplicates from the combined set
+        final_df = self._merge_duplicates(combined_df)
+
+        self.logger.info(
+            f"Ablation complete. Total prefixes after ablation and merging: {len(final_df)}"
+        )
+        self._print_detailed_stats(final_df, "ablation")
+
+        # Ensure original columns are present
+        final_cols = [col for col in original_cols if col in final_df.columns]
+        if "prefix" not in final_cols:
+            final_cols.append("prefix")
+        if "goal" not in final_cols:
+            final_cols.append("goal")
+        return final_df[list(dict.fromkeys(final_cols))]