undo list

elk/run.py

@@ -102,7 +102,7 @@ def execute(
 
  devices = select_usable_devices(self.num_gpus, min_memory=self.min_gpu_mem)
  num_devices = len(devices)
- func: Callable[[int], list[dict[str, pd.DataFrame]]] = partial(
+ func: Callable[[int], dict[str, pd.DataFrame]] = partial(
  self.apply_to_layer,
  devices=devices,
  world_size=num_devices,
@@ -113,7 +113,7 @@ def execute(
  @abstractmethod
  def apply_to_layer(
  self, layer: int, devices: list[str], world_size: int, probe_per_prompt: bool
- ) -> list[dict[str, pd.DataFrame]]:
+ ) -> dict[str, pd.DataFrame]:
  """Train or eval a reporter on a single layer."""
 
  def make_reproducible(self, seed: int):
@@ -162,7 +162,7 @@ def concatenate(self, layers):
 
  def apply_to_layers(
  self,
- func: Callable[[int], list[dict[str, pd.DataFrame]]],
+ func: Callable[[int], dict[str, pd.DataFrame]],
  num_devices: int,
  ):
  """Apply a function to each layer of the datasets in parallel
@@ -187,17 +187,17 @@ def apply_to_layers(
  df_buffers = defaultdict(list)
 
  try:
- for df_dicts in tqdm(mapper(func, layers), total=len(layers)):
- for df_dict in df_dicts:
- for k, v in df_dict.items():
- df_buffers[k].append(v)
+ for df_dict in tqdm(mapper(func, layers), total=len(layers)):
+ for k, v in df_dict.items():
+ df_buffers[k].append(v)
  finally:
  # Make sure the CSVs are written even if we crash or get interrupted
  for name, dfs in df_buffers.items():
  sortby = ["layer", "ensembling"]
  if "prompt_index" in dfs[0].columns:
  sortby.append("prompt_index")
- # TODO make the prompt index third col
+ # make the prompt index third col
+
  df = pd.concat(dfs).sort_values(by=sortby)
  out_path = self.out_dir / f"{name}.csv"
  df.round(4).to_csv(out_path, index=False)

elk/training/train.py

@@ -33,115 +33,117 @@ class MultiReporter:
  def __init__(self, reporter_results: list[ReporterTrainResult]):
  self.reporter_results: list[ReporterTrainResult] = reporter_results
  self.reporters = [r.reporter for r in reporter_results]
- train_losses = [r.train_loss for r in reporter_results]
- self.train_loss = (
- None if train_losses[0] is None else sum(train_losses) / len(train_losses)
- )
+ train_losses = [r.train_loss for r in reporter_results] if reporter_results[
+ 0].train_loss \
+ is not None else None
+ self.train_loss = sum(train_losses) / len(
+ train_losses
+ ) if train_losses is not None else None
 
  def __call__(self, h):
  credences = [r(h) for r in self.reporters]
  return torch.stack(credences).mean(dim=0)
 
 
-@dataclass
-class Elicit(Run):
- """Full specification of a reporter training run."""
+def evaluate_and_save(
+ train_loss,
+ reporter: AnyReporter | MultiReporter,
+ train_dict,
+ val_dict,
+ lr_models,
+ layer,
+):
+ row_bufs = defaultdict(list)
+ for ds_name in val_dict:
+ val_h, val_gt, val_lm_preds = val_dict[ds_name]
+ train_h, train_gt, train_lm_preds = train_dict[ds_name]
+ meta = {"dataset": ds_name, "layer": layer}
+
+ def eval_all(
+ reporter: AnyReporter | MultiReporter,
+ prompt_index: int | Literal["multi"],
+ ):
+ val_credences = reporter(val_h)
+ train_credences = reporter(train_h)
+ prompt_index = {"prompt_index": prompt_index}
+ for mode in ("none", "partial", "full"):
+ row_bufs["eval"].append(
+ {
+ **meta,
+ "ensembling": mode,
+ **evaluate_preds(val_gt, val_credences, mode).to_dict(),
+ "train_loss": train_loss,
+ **prompt_index,
+ }
+ )
 
- net: FitterConfig = subgroups(
- {"ccs": CcsConfig, "eigen": EigenFitterConfig}, default="eigen"
- )
- """Config for building the reporter network."""
+ row_bufs["train_eval"].append(
+ {
+ **meta,
+ "ensembling": mode,
+ **evaluate_preds(train_gt, train_credences, mode).to_dict(),
+ "train_loss": train_loss,
+ **prompt_index,
+ }
+ )
 
- supervised: Literal["none", "single", "inlp", "cv"] = "single"
- """Whether to train a supervised classifier, and if so, whether to use
- cross-validation. Defaults to "single", which means to train a single classifier
- on the training data. "cv" means to use cross-validation."""
+ if val_lm_preds is not None:
+ row_bufs["lm_eval"].append(
+ {
+ **meta,
+ "ensembling": mode,
+ **evaluate_preds(val_gt, val_lm_preds, mode).to_dict(),
+ **prompt_index,
+ }
+ )
 
- def evaluate_and_save(
- self,
- train_loss,
- reporter: AnyReporter | MultiReporter,
- train_dict,
- val_dict,
- lr_models,
- layer,
- ):
- row_bufs = defaultdict(list)
- for ds_name in val_dict:
- val_h, val_gt, val_lm_preds = val_dict[ds_name]
- train_h, train_gt, train_lm_preds = train_dict[ds_name]
- meta = {"dataset": ds_name, "layer": layer}
-
- def eval_all(
- reporter: AnyReporter | MultiReporter,
- prompt_index: int | Literal["multi"],
- ):
- val_credences = reporter(val_h)
- train_credences = reporter(train_h)
- prompt_index = {"prompt_index": prompt_index}
- for mode in ("none", "partial", "full"):
- row_bufs["eval"].append(
+ if train_lm_preds is not None:
+ row_bufs["train_lm_eval"].append(
  {
  **meta,
  "ensembling": mode,
- **evaluate_preds(val_gt, val_credences, mode).to_dict(),
- "train_loss": train_loss,
+ **evaluate_preds(
+ train_gt, train_lm_preds, mode
+ ).to_dict(),
  **prompt_index,
  }
  )
 
- row_bufs["train_eval"].append(
+ for i, model in enumerate(lr_models):
+ row_bufs["lr_eval"].append(
  {
  **meta,
  "ensembling": mode,
- **evaluate_preds(train_gt, train_credences, mode).to_dict(),
- "train_loss": train_loss,
+ "inlp_iter": i,
+ **evaluate_preds(val_gt, model(val_h), mode).to_dict(),
  **prompt_index,
  }
  )
 
- if val_lm_preds is not None:
- row_bufs["lm_eval"].append(
- {
- **meta,
- "ensembling": mode,
- **evaluate_preds(val_gt, val_lm_preds, mode).to_dict(),
- **prompt_index,
- }
- )
-
- if train_lm_preds is not None:
- row_bufs["train_lm_eval"].append(
- {
- **meta,
- "ensembling": mode,
- **evaluate_preds(
- train_gt, train_lm_preds, mode
- ).to_dict(),
- **prompt_index,
- }
- )
-
- for i, model in enumerate(lr_models):
- row_bufs["lr_eval"].append(
- {
- **meta,
- "ensembling": mode,
- "inlp_iter": i,
- **evaluate_preds(val_gt, model(val_h), mode).to_dict(),
- **prompt_index,
- }
- )
-
- if isinstance(reporter, MultiReporter):
- for prompt_index, reporter_result in enumerate(
- reporter.reporter_results
- ):
- eval_all(reporter_result.reporter, prompt_index)
-
- eval_all(reporter, "multi")
-
- return {k: pd.DataFrame(v) for k, v in row_bufs.items()}
+ if isinstance(reporter, MultiReporter):
+ for prompt_index, reporter_result in enumerate(
+ reporter.reporter_results
+ ):
+ eval_all(reporter_result.reporter, prompt_index)
+
+ eval_all(reporter, "multi")
+
+ return {k: pd.DataFrame(v) for k, v in row_bufs.items()}
+
+
+@dataclass
+class Elicit(Run):
+ """Full specification of a reporter training run."""
+
+ net: FitterConfig = subgroups(
+ {"ccs": CcsConfig, "eigen": EigenFitterConfig}, default="eigen"
+ )
+ """Config for building the reporter network."""
+
+ supervised: Literal["none", "single", "inlp", "cv"] = "single"
+ """Whether to train a supervised classifier, and if so, whether to use
+ cross-validation. Defaults to "single", which means to train a single classifier
+ on the training data. "cv" means to use cross-validation."""
 
  def create_models_dir(self, out_dir: Path):
  lr_dir = None
@@ -258,7 +260,7 @@ def apply_to_layer(
  devices: list[str],
  world_size: int,
  probe_per_prompt: bool,
- ) -> list[dict[str, pd.DataFrame]]:
+ ) -> dict[str, pd.DataFrame]:
  """Train a single reporter on a single layer."""
  assert self.out_dir is not None # TODO this is really annoying, why can it be
  # None?
@@ -275,14 +277,13 @@ def apply_to_layer(
  # TODO is this even needed
  # reporter_dir, lr_dir = self.create_models_dir(assert_type(Path, self.out_dir))
 
- probe_per_prompt = True
  if probe_per_prompt:
  train_dicts = [
  {
  ds_name: (
- train_h[:, i : i + 1, ...],
+ train_h[:, i: i + 1, ...],
  train_gt,
- lm_preds[:, i : i + 1, ...] if lm_preds is not None else None,
+ lm_preds[:, i: i + 1, ...] if lm_preds is not None else None,
  )
  }
  for ds_name, (train_h, _, lm_preds) in train_dict.items()
@@ -304,16 +305,14 @@ def apply_to_layer(
  multi_reporter = MultiReporter(results)
  train_loss = multi_reporter.train_loss
 
- return [
- self.evaluate_and_save(
- train_loss,
- multi_reporter,
- train_dict,
- val_dict,
- lr_models, # TODO I don't care about this right now but
- layer,
- )
- ]
+ return evaluate_and_save(
+ train_loss,
+ multi_reporter,
+ train_dict,
+ val_dict,
+ lr_models, # TODO I don't care about this right now but
+ layer,
+ )
  else:
  reporter_train_result = self.train_and_save_reporter(
  device, layer, self.out_dir / "reporters", train_dict
@@ -326,8 +325,6 @@ def apply_to_layer(
  train_dict, device, layer, self.out_dir / "lr_models"
  )
 
- return [
- self.evaluate_and_save(
- train_loss, reporter, train_dict, val_dict, lr_models, layer
- )
- ]
+ return evaluate_and_save(
+ train_loss, reporter, train_dict, val_dict, lr_models, layer
+ )