minor edits

elk/extraction/prompt_dataset.py

@@ -120,15 +120,15 @@ def __init__(
  # This allows you to just set split="train" and split="test" for any dataset
  # and not worry about train-test leakage.
 
- # split_name, *others = ds_dict.keys()
- # if not others:
- #  print("Creating a 75/25 train-test split...")
+ split_name, *others = ds_dict.keys()
+ if not others:
+ print("Creating a 75/25 train-test split...")
 
- #  # Don't shuffle now because we're going to shuffle later
- #  ds_dict = ds_dict[split_name].train_test_split(
- #  seed=cfg.seed, shuffle=False, stratify_by_column=cfg.label_column
- #  )
- #  assert isinstance(ds_dict, DatasetDict)
+ # Don't shuffle now because we're going to shuffle later
+ ds_dict = ds_dict[split_name].train_test_split(
+ seed=cfg.seed, shuffle=False, stratify_by_column=cfg.label_column
+ )
+ assert isinstance(ds_dict, DatasetDict)
 
  # The 'active' split is the one that gets queried by __getitem__
  self.active_split = ds_dict[split]
@@ -254,7 +254,7 @@ def num_classes(self) -> int:
  return self.active_split.features[self.label_column].num_classes
 
 
-class Interleaved_Datasets(TorchDataset):
+class InterleavedDatasets(TorchDataset):
  def __init__(
  self,
  datasets: list[PromptDataset],

elk/training/train.py

@@ -52,14 +52,14 @@ def train_reporter(
  cfg: RunConfig,
  dataset: DatasetDict,
  out_dir: Path,
- layer: Union[int, list[int]],
+ layer: list[int],
  devices: list[str],
  world_size: int = 1,
 ):
  """Train a single reporter on a single layer, or a list of layers."""
 
  # Reproducibility
- seed = cfg.net.seed + layer if isinstance(layer, int) else layer[0]
+ seed = cfg.net.seed + layer[0]
  np.random.seed(seed)
  random.seed(seed)
  torch.manual_seed(seed)
@@ -180,16 +180,17 @@ def train(cfg: RunConfig, out_dir: Optional[Path] = None):
  if not cfg.skip_baseline:
  cols += ["lr_auroc", "lr_acc"]
 
+ # Create subsets of layers to train reporters on
  layers = [
- int(feat[len("hidden_") :])
+ [int(feat[len("hidden_") :])]
  for feat in ds["train"].features
  if feat.startswith("hidden_")
  ]
 
  # concatenate hidden states from a previous layer, if told to
  if cfg.concatenate_layers > 0:
  for i in range(cfg.concatenate_layers, len(layers)):
- layers[i] = [layers[i], layers[i] - cfg.concatenate_layers]
+ layers[i] = layers[i] + [layers[i][0] - cfg.concatenate_layers]
 
  # Train reporters for each layer in parallel
  with mp.Pool(num_devices) as pool, open(out_dir / "eval.csv", "w") as f: