From 294b25ca1f531f2fd97333292dd30315a009e225 Mon Sep 17 00:00:00 2001
From: Adam Karvonen <adam.karvonen@gmail.com>
Date: Sat, 18 Jan 2025 19:24:29 +0000
Subject: [PATCH] Check that unlearning data exists before running unlearning
 eval

---
 .../custom_saes/run_all_evals_custom_saes.py  |  4 +++
 .../run_all_evals_dictionary_learning_saes.py | 25 ++++++++++++-------
 .../misc_notebooks/export_results.ipynb       | 17 ++++++++++---
 3 files changed, 34 insertions(+), 12 deletions(-)

diff --git a/sae_bench/custom_saes/run_all_evals_custom_saes.py b/sae_bench/custom_saes/run_all_evals_custom_saes.py
index 7c498d8..fbf787c 100644
--- a/sae_bench/custom_saes/run_all_evals_custom_saes.py
+++ b/sae_bench/custom_saes/run_all_evals_custom_saes.py
@@ -228,6 +228,10 @@ def run_evals(
     else:
         api_key = None
 
+    if "unlearning" in eval_types:
+        if not os.path.exists("./sae_bench/evals/unlearning/data/bio-forget-corpus.jsonl"):
+            raise Exception("Please download bio-forget-corpus.jsonl for unlearning evaluation")
+
     # If evaluating multiple SAEs on the same layer, set save_activations to True
     # This will require at least 100GB of disk space
     save_activations = False
diff --git a/sae_bench/custom_saes/run_all_evals_dictionary_learning_saes.py b/sae_bench/custom_saes/run_all_evals_dictionary_learning_saes.py
index b128c50..5ea2185 100644
--- a/sae_bench/custom_saes/run_all_evals_dictionary_learning_saes.py
+++ b/sae_bench/custom_saes/run_all_evals_dictionary_learning_saes.py
@@ -333,6 +333,9 @@ def run_evals(
     device = general_utils.setup_environment()
 
     # Select your eval types here.
+    # Note: Unlearning is not recommended for models with < 2B parameters and we recommend an instruct tuned model
+    # Unlearning will also require requesting permission for the WMDP dataset (see unlearning/README.md)
+    # Absorption not recommended for models < 2B parameters
     eval_types = [
         "absorption",
         "core",
@@ -343,6 +346,19 @@ def run_evals(
         # "unlearning",
     ]
 
+    if "autointerp" in eval_types:
+        try:
+            with open("openai_api_key.txt") as f:
+                api_key = f.read().strip()
+        except FileNotFoundError:
+            raise Exception("Please create openai_api_key.txt with your API key")
+    else:
+        api_key = None
+
+    if "unlearning" in eval_types:
+        if not os.path.exists("./sae_bench/evals/unlearning/data/bio-forget-corpus.jsonl"):
+            raise Exception("Please download bio-forget-corpus.jsonl for unlearning evaluation")
+
     repos = [
         (
             "adamkarvonen/saebench_pythia-160m-deduped_width-2pow14_date-0104",
@@ -369,15 +385,6 @@ def run_evals(
         # Unlearning will also require requesting permission for the WMDP dataset (see unlearning/README.md)
         # Absorption not recommended for models < 2B parameters
 
-        if "autointerp" in eval_types:
-            try:
-                with open("openai_api_key.txt") as f:
-                    api_key = f.read().strip()
-            except FileNotFoundError:
-                raise Exception("Please create openai_api_key.txt with your API key")
-        else:
-            api_key = None
-
         run_evals(
             repo_id=repo_id,
             model_name=model_name,
diff --git a/sae_bench/sae_bench_utils/misc_notebooks/export_results.ipynb b/sae_bench/sae_bench_utils/misc_notebooks/export_results.ipynb
index 716fb41..2dafe46 100644
--- a/sae_bench/sae_bench_utils/misc_notebooks/export_results.ipynb
+++ b/sae_bench/sae_bench_utils/misc_notebooks/export_results.ipynb
@@ -96,6 +96,20 @@
     "            print(f\"Folder {folder} does not exist. Skipping.\")\n",
     "            continue\n",
     "\n",
+    "        # Get all json files in the current directory\n",
+    "        files = [f for f in os.listdir(\".\") if f.endswith(\".json\")]\n",
+    "        if len(files) == 0:\n",
+    "            print(f\"No json files found in {folder}. Skipping.\")\n",
+    "            continue\n",
+    "\n",
+    "        num_files_with_prefix = 0\n",
+    "\n",
+    "        for prefix in prefixes:\n",
+    "            num_files_with_prefix += len([f for f in files if f.startswith(prefix)])\n",
+    "\n",
+    "        if num_files_with_prefix == 0:\n",
+    "            raise ValueError(f\"No files with prefix {prefixes} found in {folder}.\")\n",
+    "\n",
     "        os.chdir(folder)  # Change to the target directory\n",
     "        print(f\"Processing folder: {folder}\")\n",
     "\n",
@@ -104,9 +118,6 @@
     "            if not os.path.exists(prefix):\n",
     "                os.makedirs(prefix)\n",
     "\n",
-    "        # Get all json files in the current directory\n",
-    "        files = [f for f in os.listdir(\".\") if f.endswith(\".json\")]\n",
-    "\n",
     "        # Move files to appropriate folders\n",
     "        for file in files:\n",
     "            for prefix in prefixes:\n",