Merge with EleutherELK

elk/training/sweep.py

@@ -1,12 +1,11 @@
-from dataclasses import InitVar, dataclass, replace, field
+from dataclasses import InitVar, dataclass, field, replace
 
 import numpy as np
 import torch
 from datasets import get_dataset_config_info
 from transformers import AutoConfig
 
 from ..evaluation import Eval
-from ..extraction import Extract
 from ..files import memorably_named_dir, sweeps_dir
 from ..plotting.visualize import visualize_sweep
 from ..training.eigen_reporter import EigenFitterConfig
@@ -53,13 +52,7 @@ class Sweep:
     name: str | None = None
 
     # A bit of a hack to add all the command line arguments from Elicit
-    run_template: Elicit = field(default_factory=lambda: Elicit(
-            data=Extract(
-                model="<placeholder>",
-                datasets=("<placeholder>",),
-            )
-        )
-    )
+    run_template: Elicit = field(default_factory=Elicit.default)
 
     def __post_init__(self, add_pooled: bool):
         if not self.datasets:

elk/training/train.py

@@ -11,6 +11,7 @@
 from simple_parsing import subgroups
 from simple_parsing.helpers.serialization import save
 
+from ..extraction import Extract
 from ..metrics import evaluate_preds, to_one_hot
 from ..run import Run
 from ..training.supervised import train_supervised
@@ -34,6 +35,15 @@ 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."""
 
+    @staticmethod
+    def default():
+        return Elicit(
+            data=Extract(
+                model="<placeholder>",
+                datasets=("<placeholder>",),
+            )
+        )
+
     def create_models_dir(self, out_dir: Path):
         lr_dir = None
         lr_dir = out_dir / "lr_models"

pyproject.toml

@@ -12,6 +12,8 @@ license = { text = "MIT License" }
 dependencies = [
     # Allows us to use device_map in from_pretrained. Also needed for 8bit
     "accelerate",
+    # Already a dependency of datasets, but newer versions introduce breaking changes
+    "pyarrow==12.0.0",
     # For pseudolabel and prompt normalization. We're picky about the version because
     # the package isn't guaranteed to be stable yet.
     "concept-erasure==0.1.0",