fix ruff errors

Author: robbiemccorkell

attribution/experiment_logger.py

@@ -2,13 +2,14 @@
 import pickle
 import time
 from typing import Any, Literal, Optional
-from IPython.display import HTML
+
+import matplotlib.colors as mcolors
+import matplotlib.pyplot as plt
 import pandas as pd
 from IPython.core.getipython import get_ipython
+from IPython.display import HTML
 
 from .token_perturbation import PerturbationStrategy, PerturbedLLMInput, combine_unit
-import matplotlib.pyplot as plt
-import matplotlib.colors as mcolors
 
 
 class ExperimentLogger:
@@ -177,23 +178,26 @@ def log_perturbation(
     def clean_tokens(self, tokens):
         return [token.replace("Ġ", " ") for token in tokens]
 
-
     def score_to_color(self, score, vmin=-1, vmax=1):
         norm = mcolors.Normalize(vmin=vmin, vmax=vmax)
         cmap = plt.cm.coolwarm
         rgba_color = cmap(norm(score))
         color_hex = mcolors.to_hex(rgba_color)
         return color_hex
 
-    def print_text_total_attribution(self, exp_id: Optional[int] = None, score_agg: Literal["mean", "sum", "last"] = "mean"):
-
+    def print_text_total_attribution(
+        self, exp_id: Optional[int] = None, score_agg: Literal["mean", "sum", "last"] = "mean"
+    ):
         if exp_id == -1:
             exp_id = self.df_experiments["exp_id"].max()
 
-        token_attrs_df = self.df_input_token_attribution.groupby(
-            ["exp_id", "attribution_strategy"]) if exp_id is None else self.df_input_token_attribution[self.df_input_token_attribution["exp_id"] == exp_id].groupby(
-            ["exp_id", "attribution_strategy"])
-        
+        token_attrs_df = (
+            self.df_input_token_attribution.groupby(["exp_id", "attribution_strategy"])
+            if exp_id is None
+            else self.df_input_token_attribution[
+                self.df_input_token_attribution["exp_id"] == exp_id
+            ].groupby(["exp_id", "attribution_strategy"])
+        )
 
         for (exp_id, attr_strat), exp_data in token_attrs_df:
             if exp_data["input_token_pos"].duplicated().any():
@@ -221,40 +225,40 @@ def print_text_total_attribution(self, exp_id: Optional[int] = None, score_agg:
                 color = self.score_to_color(score, vmin, vmax)
                 html_str += f'<span style="text-decoration: underline; text-decoration-color: {color}; text-decoration-thickness: 4px;">{token}</span>'
 
-
-            html_str += f' -> <b>{output}</b>'     
-            html_str += '</div>'
+            html_str += f" -> <b>{output}</b>"
+            html_str += "</div>"
 
             # Display
-            
+
             if get_ipython() and "IPKernelApp" in get_ipython().config:
                 from IPython.display import display
+
                 display(HTML(html_str))
             else:
                 self.pretty_print(df)
 
-
     def print_text_attribution_matrix(self, exp_id: int = -1):
-
         if exp_id == -1:
             exp_id = self.df_experiments["exp_id"].max()
 
         matrix = self.get_attribution_matrix(exp_id)
 
-        input_tokens = [' '.join(x.split(' ')[:-1]) for x in matrix.index]
+        input_tokens = [" ".join(x.split(" ")[:-1]) for x in matrix.index]
         token_dict = {f"token_{i+1}": t for i, t in enumerate(input_tokens)}
 
         for oi, output_token in enumerate(matrix.columns):
-            prev_output_str = ''.join([' '.join(ot.split(' ')[:-1]) for ot in matrix.columns[:oi]])
-            following_output_str = ''.join([' '.join(ot.split(' ')[:-1]) for ot in matrix.columns[oi+1:]])
+            prev_output_str = "".join([" ".join(ot.split(" ")[:-1]) for ot in matrix.columns[:oi]])
+            following_output_str = "".join(
+                [" ".join(ot.split(" ")[:-1]) for ot in matrix.columns[oi + 1 :]]
+            )
             attr_scores = matrix[output_token].tolist()
             vmax = max(abs(min(attr_scores)), abs(max(attr_scores)))
             vmin = -vmax
 
             score_dict = {f"token_{i+1}": score for i, score in enumerate(attr_scores)}
 
             df = pd.DataFrame([token_dict, score_dict], index=["token", "attr_score"])
-            
+
             # Generating HTML
             html_str = '<div style="font-family: monospace;">'
             for col in df.columns:
@@ -263,27 +267,33 @@ def print_text_attribution_matrix(self, exp_id: int = -1):
                 color = self.score_to_color(score, vmin, vmax)
                 html_str += f'<span style="text-decoration: underline; text-decoration-color: {color}; text-decoration-thickness: 4px;">{token}</span>'
 
-            clean_output_token = ' '.join(output_token.split(' ')[:-1])
-            html_str += f' -> {prev_output_str}<b>{clean_output_token}</b>{following_output_str}'     
-            html_str += '</div>'
+            clean_output_token = " ".join(output_token.split(" ")[:-1])
+            html_str += f" -> {prev_output_str}<b>{clean_output_token}</b>{following_output_str}"
+            html_str += "</div>"
 
             # Display
-            
+
             if get_ipython() and "IPKernelApp" in get_ipython().config:
                 from IPython.display import display
+
                 display(HTML(html_str))
             else:
                 self.pretty_print(df)
 
-
-    def print_total_attribution(self, exp_id: Optional[int] = None, score_agg: Literal["mean", "last"] = "mean"):
+    def print_total_attribution(
+        self, exp_id: Optional[int] = None, score_agg: Literal["mean", "last"] = "mean"
+    ):
         totals = []
         if exp_id == -1:
             exp_id = self.df_experiments["exp_id"].max()
 
-        token_attrs_df = self.df_input_token_attribution.groupby(
-            ["exp_id", "attribution_strategy"]) if exp_id is None else self.df_input_token_attribution[self.df_input_token_attribution["exp_id"] == exp_id].groupby(
-            ["exp_id", "attribution_strategy"])
+        token_attrs_df = (
+            self.df_input_token_attribution.groupby(["exp_id", "attribution_strategy"])
+            if exp_id is None
+            else self.df_input_token_attribution[
+                self.df_input_token_attribution["exp_id"] == exp_id
+            ].groupby(["exp_id", "attribution_strategy"])
+        )
 
         for (exp_id, attr_strat), exp_data in token_attrs_df:
             if exp_data["input_token_pos"].duplicated().any():
@@ -324,11 +334,18 @@ def print_attribution_matrix(
     ):
         if exp_id == -1:
             exp_id = self.df_experiments["exp_id"].max()
-        matrix = self.get_attribution_matrix(exp_id, attribution_strategy, show_debug_cols, score_agg)
+        matrix = self.get_attribution_matrix(
+            exp_id, attribution_strategy, show_debug_cols, score_agg
+        )
 
         if get_ipython() and "IPKernelApp" in get_ipython().config:
             from IPython.display import display
-            display(matrix.style.background_gradient(cmap="coolwarm", vmin=-1, vmax=1).set_properties(**{"white-space": "pre-wrap"}))
+
+            display(
+                matrix.style.background_gradient(cmap="coolwarm", vmin=-1, vmax=1).set_properties(
+                    **{"white-space": "pre-wrap"}
+                )
+            )
         else:
             self.pretty_print(matrix)
 
@@ -343,13 +360,14 @@ def get_attribution_matrix(
             exp_id = self.df_experiments["exp_id"].max()
 
         if attribution_strategy is None:
-            strategies = self.df_token_attribution_matrix[(self.df_token_attribution_matrix["exp_id"] == exp_id)]["attribution_strategy"].unique()
+            strategies = self.df_token_attribution_matrix[
+                (self.df_token_attribution_matrix["exp_id"] == exp_id)
+            ]["attribution_strategy"].unique()
         else:
             strategies = [attribution_strategy]
 
         matrices = []
         for attribution_strategy in strategies:
-                
             # Filter the data for the specific experiment and attribution strategy
             exp_data = self.df_token_attribution_matrix[
                 (self.df_token_attribution_matrix["exp_id"] == exp_id)
@@ -401,14 +419,15 @@ def get_attribution_matrix(
                 )
                 additional_columns = additional_columns[["perturbed_input", "perturbed_output"]]
                 matrix = matrix.join(additional_columns)
-            
+
             matrices.append(matrix)
         return pd.concat(matrices)
 
     def pretty_print(self, df: pd.DataFrame):
         # Check if code is running in Jupyter notebook
         if get_ipython() and "IPKernelApp" in get_ipython().config:
             from IPython.display import display
+
             display(df.style.set_properties(**{"white-space": "pre-wrap"}))
         else:
             print(df.to_string())

examples/example_PIZZA.ipynb

@@ -15,13 +15,12 @@
       "metadata": {},
       "outputs": [],
       "source": [
+        "import warnings\n",
+        "\n",
         "# Set your open ai API key\n",
         "# BEWARE: This will cost you API credits!\n",
         "YOUR_OPENAI_API_KEY = \"your-api-key\"\n",
         "\n",
-        "\n",
-        "import warnings\n",
-        "\n",
         "# Suppress annoying FutureWarning from huggingface_hub\n",
         "warnings.filterwarnings(\"ignore\", category=FutureWarning, module=\"huggingface_hub\")"
       ]