fix save_logprobs

elk/__init__.py

@@ -1,5 +1,9 @@
+from .evaluation import Eval
 from .extraction import Extract
+from .training.train import Elicit
 
 __all__ = [
     "Extract",
+    "Elicit",
+    "Eval",
 ]

elk/evaluation/evaluate.py

@@ -73,6 +73,9 @@ def apply_to_layer(
                         }
                     )
 
+                if self.save_logprobs:
+                    out_logprobs[ds_name]["lr"][mode] = dict()
+
                 for i, model in enumerate(lr_models):
                     model.eval()
                     val_log_odds = model(val_data.hiddens)

elk/promptsource/templates/_default/templates.yaml

@@ -2,7 +2,7 @@ dataset: None
 templates:
   7eab7254-bd71-4b1d-9f8a-0fc7110f8371: !Template
     answer_choices: False ||| True
-    id: 7eab7254-bd41-4b1d-9f8a-0fc7110f8371
+    id: 7eab7254-bd71-4b1d-9f8a-0fc7110f8371
     jinja: "{{ statement }}"
     metadata: !TemplateMetadata
       choices_in_prompt: true

elk/promptsource/templates/_no_suffix/templates.yaml

@@ -0,0 +1,15 @@
+dataset: None
+templates:
+  8eab7252-bd71-4b2d-9f8a-0fc7260f8371: !Template
+    answer_choices: False ||| True
+    id: 8eab7252-bd71-4b2d-9f8a-0fc7260f8371
+    jinja: "{{ statement }}"
+    metadata: !TemplateMetadata
+      choices_in_prompt: true
+      languages:
+      - en
+      metrics:
+      - Accuracy
+      original_task: true
+    name: _no_suffix
+    suffix: ""

elk/run.py

@@ -68,7 +68,7 @@ class Run(ABC, Serializable):
     concatenated_layer_offset: int = 0
     debug: bool = False
     num_gpus: int = -1
-    out_dir: Path | None = None
+    min_gpu_mem: int = 0
     disable_cache: bool = field(default=False, to_dict=False)
 
     def execute(
@@ -113,7 +113,7 @@ def execute(
                 meta_f,
             )
 
-        devices = select_usable_devices(self.num_gpus)
+        devices = select_usable_devices(self.num_gpus, min_memory=self.min_gpu_mem)
         num_devices = len(devices)
         func: Callable[[int], tuple[dict[str, pd.DataFrame], dict]] = partial(
             self.apply_to_layer, devices=devices, world_size=num_devices
@@ -225,18 +225,17 @@ def apply_to_layers(
                 if self.save_logprobs:
                     save_dict = defaultdict(dict)
                     for ds_name, logprobs_dict in logprobs_dicts.items():
+                        save_dict[ds_name]["row_ids"] = logprobs_dict[layers[0]][
+                            "row_ids"
+                        ]
                         save_dict[ds_name]["texts"] = logprobs_dict[layers[0]]["texts"]
                         save_dict[ds_name]["labels"] = logprobs_dict[layers[0]][
                             "labels"
                         ]
                         save_dict[ds_name]["lm"] = logprobs_dict[layers[0]]["lm"]
-                        save_dict[ds_name]["reporter"] = dict()
                         save_dict[ds_name]["lr"] = dict()
                         for layer, logprobs_dict_by_mode in logprobs_dict.items():
-                            save_dict[ds_name]["reporter"][
-                                layer
-                            ] = logprobs_dict_by_mode["reporter"]
                             save_dict[ds_name]["lr"][layer] = logprobs_dict_by_mode[
                                 "lr"
                             ]
-                    torch.save(save_dict, self.out_dir / "logprobs.pt")
+                    torch.save(dict(save_dict), self.out_dir / "logprobs.pt")

elk/training/supervised.py

@@ -24,7 +24,7 @@ def train_supervised(
         (n, v, d) = train_data.hiddens.shape
         train_h = rearrange(train_data.hiddens, "n v d -> (n v) d")
 
-        if erase_paraphrases:
+        if erase_paraphrases and v > 1:
             if leace is None:
                 leace = LeaceFitter(
                     d,

elk/training/train.py

@@ -26,7 +26,7 @@ class Elicit(Run):
     cross-validation. Defaults to "single", which means to train a single classifier
     on the training data. "cv" means to use cross-validation."""
 
-    erase_paraphrases: bool = True
+    erase_paraphrases: bool = False
     """Whether to use LEACE to erase the paraphrase dimensions before training the
     classifier."""
 
@@ -41,7 +41,7 @@ def default():
                 datasets=("<placeholder>",),
             )
         )
-    
+
     def create_models_dir(self, out_dir: Path):
         lr_dir = out_dir / "lr_models"
 
@@ -115,6 +115,9 @@ def apply_to_layer(
                         }
                     )
 
+                if self.save_logprobs:
+                    out_logprobs[ds_name]["lr"][mode] = dict()
+
                 for i, model in enumerate(lr_models):
                     model.eval()
                     val_log_odds = model(val.hiddens)