Fix

etowahadams · Oct 16, 2024 · cd1411d · cd1411d
1 parent 8ce29b1
commit cd1411d
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 13 deletions.
diff --git a/plm_interpretability/logistic_regression_probe/README.md b/plm_interpretability/logistic_regression_probe/README.md
@@ -3,11 +3,27 @@
 ## Single latent
 
 ```bash
-logistic_regression_probe single_latent --sae-checkpoint plm_interpretability/checkpoints/l24_plm1280_sae4096_k128_100k.pt --sae-dim 4096 --plm-dim 1280 --plm-layer 24 --swissprot-tsv plm_interpretability/logistic_regression_probe/data/swissprot.tsv --output-dir plm_interpretability/logistic_regression_probe/results
+logistic_regression_probe single-latent \
+--sae-checkpoint plm_interpretability/checkpoints/l24_plm1280_sae4096_k128_100k.pt \
+--sae-dim 4096 \
+--plm-dim 1280 \
+--plm-layer 24 \
+--swissprot-tsv plm_interpretability/logistic_regression_probe/data/swissprot.tsv \
+--output-dir plm_interpretability/logistic_regression_probe/results \
+--max-seqs-per-task 5 \
+--annotation-names "DNA binding"
 ```
 
 ## All latents
 
 ```bash
-logistic_regression_probe all_latents --sae-checkpoint plm_interpretability/checkpoints/l24_plm1280_sae4096_k128_100k.pt --sae-dim 4096 --plm-dim 1280 --plm-layer 24 --swissprot-tsv plm_interpretability/logistic_regression_probe/data/swissprot.tsv --output-dir plm_interpretability/logistic_regression_probe/results
+logistic_regression_probe all-latents \
+--sae-checkpoint plm_interpretability/checkpoints/l24_plm1280_sae4096_k128_100k.pt \
+--sae-dim 4096 \
+--plm-dim 1280 \
+--plm-layer 24 \
+--swissprot-tsv plm_interpretability/logistic_regression_probe/data/swissprot.tsv \
+--output-file plm_interpretability/logistic_regression_probe/results/all_latents.csv \
+--max-seqs-per-task 5 \
+--annotation-names "DNA binding"
 ```
diff --git a/plm_interpretability/logistic_regression_probe/all_latents.py b/plm_interpretability/logistic_regression_probe/all_latents.py
@@ -42,7 +42,7 @@
 )
 @click.option(
     "--output-file",
-    type=click.File(),
+    type=click.Path(),
     required=True,
     help="Path to the output file",
 )
@@ -131,7 +131,3 @@ def all_latents(
 
             del seq_to_annotation_entries, examples, train_examples, test_examples
             gc.collect()
-
-
-if __name__ == "__main__":
-    all_latents()
diff --git a/plm_interpretability/logistic_regression_probe/single_latent.py b/plm_interpretability/logistic_regression_probe/single_latent.py
@@ -224,7 +224,3 @@ def single_latent(
 
             del seq_to_annotation_entries, examples, train_examples, test_examples
             gc.collect()
-
-
-if __name__ == "__main__":
-    single_latent()
diff --git a/plm_interpretability/logistic_regression_probe/utils.py b/plm_interpretability/logistic_regression_probe/utils.py
@@ -10,6 +10,8 @@
 from plm_interpretability.sae_model import SparseAutoencoder
 from plm_interpretability.utils import get_layer_activations, parse_swissprot_annotation
 
+MAX_SEQ_LEN = 1000
+
 
 def get_sae_acts(
     seq: str,
@@ -56,12 +58,12 @@ def get_annotation_entries_for_class(
             # The note field is sometimes like "Homeobox", "Homeobox 1", etc.,
             # so use string `in` to check.
             entries = [e for e in entries if class_name in e.get("note", "")]
-        if len(entries) > 0 and len(seq) < 2000:
+        if len(entries) > 0 and len(seq) < MAX_SEQ_LEN:
             seq_to_annotation_entries[seq] = entries
             seq_lengths.append(len(seq))
 
     logger.info(
-        f"Found {len(seq_to_annotation_entries)} sequences with class {class_name}."
+        f"Found {len(seq_to_annotation_entries)} sequences with class {class_name}. "
         f"Mean sequence length: {np.mean(seq_lengths):.2f}."
     )