From 7d0e3961903fbc559bc6f7e92b497d87b5a34244 Mon Sep 17 00:00:00 2001
From: Johnny Lin <hijohnnylin@gmail.com>
Date: Mon, 22 Apr 2024 16:26:19 -0700
Subject: [PATCH] skip batch file if it already exists

---
 sae_lens/analysis/neuronpedia_runner.py | 11 ++++++++++-
 tutorials/neuronpedia/neuronpedia.py    |  7 -------
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/sae_lens/analysis/neuronpedia_runner.py b/sae_lens/analysis/neuronpedia_runner.py
index ff580069..95b09cb7 100644
--- a/sae_lens/analysis/neuronpedia_runner.py
+++ b/sae_lens/analysis/neuronpedia_runner.py
@@ -238,6 +238,13 @@ def run(self):
                     # print(f"Skipping batch - it's after end_batch: {feature_batch_count}")
                     continue
 
+                output_file = f"{self.outputs_dir}/batch-{feature_batch_count}.json"
+                # if output_file exists, skip
+                if os.path.isfile(output_file):
+                    logline = f"\n++++++++++ Skipping Batch #{feature_batch_count} output. File exists: {output_file} ++++++++++\n"
+                    print(logline)
+                    continue
+
                 print(f"========== Running Batch #{feature_batch_count} ==========")
 
                 layout = SaeVisLayoutConfig(
@@ -422,9 +429,11 @@ def run(self):
                 json_object = json.dumps(to_write, cls=NpEncoder)
 
                 with open(
-                    f"{self.outputs_dir}/batch-{feature_batch_count}.json",
+                    output_file,
                     "w",
                 ) as f:
                     f.write(json_object)
 
+                logline = f"\n========== Completed Batch #{feature_batch_count} output: {output_file} ==========\n"
+
         return
diff --git a/tutorials/neuronpedia/neuronpedia.py b/tutorials/neuronpedia/neuronpedia.py
index aa36a6b6..e855ff28 100755
--- a/tutorials/neuronpedia/neuronpedia.py
+++ b/tutorials/neuronpedia/neuronpedia.py
@@ -147,13 +147,6 @@ def generate(
         print(f"Error: Output directory {outputs_dir.as_posix()} exists and is a file.")
         raise typer.Abort()
     outputs_dir.mkdir(parents=True, exist_ok=True)
-    # Check if output_dir has any files starting with "batch_"
-    batch_files = list(outputs_dir.glob("batch-*.json"))
-    if len(batch_files) > 0 and resume_from_batch == 1:
-        print(
-            f"Error: Output directory {outputs_dir.as_posix()} has existing batch files. This is only allowed if you are resuming from a batch. Please delete or move the existing batch-*.json files."
-        )
-        raise typer.Abort()
 
     sparsity = load_sparsity(sae_path_string)
     # convert sparsity to logged sparsity if it's not