train probe per prompt

elk/evaluation/evaluate.py

@@ -1,6 +1,7 @@
 from collections import defaultdict
 from dataclasses import dataclass
 from pathlib import Path
+from typing import Literal
 
 import pandas as pd
 import torch
@@ -9,6 +10,7 @@
 from ..files import elk_reporter_dir
 from ..metrics import evaluate_preds
 from ..run import Run
+from ..training.multi_reporter import MultiReporter, SingleReporter
 from ..utils import Color
 
 
@@ -30,47 +32,78 @@ def execute(self, highlight_color: Color = "cyan"):
 
  @torch.inference_mode()
  def apply_to_layer(
- self, layer: int, devices: list[str], world_size: int
+ self, layer: int, devices: list[str], world_size: int, probe_per_prompt: bool
  ) -> dict[str, pd.DataFrame]:
  """Evaluate a single reporter on a single layer."""
  device = self.get_device(devices, world_size)
  val_output = self.prepare_data(device, layer, "val")
 
  experiment_dir = elk_reporter_dir() / self.source
 
- reporter_path = experiment_dir / "reporters" / f"layer_{layer}.pt"
- reporter = torch.load(reporter_path, map_location=device)
+ def load_reporter() -> SingleReporter | MultiReporter:
+ # check if experiment_dir / "reporters" has .pt files
+ first = next((experiment_dir / "reporters").iterdir())
+ if not first.suffix == ".pt":
+ return MultiReporter.load(
+ experiment_dir / "reporters", layer, device=device
+ )
+ else:
+ path = experiment_dir / "reporters" / f"layer_{layer}.pt"
+ return torch.load(path, map_location=device)
+
+ reporter = load_reporter()
 
  row_bufs = defaultdict(list)
- for ds_name, (val_h, val_gt, _) in val_output.items():
- meta = {"dataset": ds_name, "layer": layer}
-
- val_credences = reporter(val_h)
- for mode in ("none", "partial", "full"):
- row_bufs["eval"].append(
- {
- **meta,
- "ensembling": mode,
- **evaluate_preds(val_gt, val_credences, mode).to_dict(),
- }
- )
 
- lr_dir = experiment_dir / "lr_models"
- if not self.skip_supervised and lr_dir.exists():
- with open(lr_dir / f"layer_{layer}.pt", "rb") as f:
- lr_models = torch.load(f, map_location=device)
- if not isinstance(lr_models, list): # backward compatibility
- lr_models = [lr_models]
-
- for i, model in enumerate(lr_models):
- model.eval()
- row_bufs["lr_eval"].append(
- {
- "ensembling": mode,
- "inlp_iter": i,
- **meta,
- **evaluate_preds(val_gt, model(val_h), mode).to_dict(),
- }
- )
+ def eval_all(
+ reporter: SingleReporter | MultiReporter,
+ prompt_index: int | Literal["multi"] | None = None,
+ i: int = 0,
+ ):
+ prompt_index_dict = (
+ {"prompt_index": prompt_index} if prompt_index is not None else {}
+ )
+ for ds_name, (val_h, val_gt, _) in val_output.items():
+ meta = {"dataset": ds_name, "layer": layer}
+ val_credences = reporter(val_h[:, [i], :, :])
+
+ for mode in ("none", "partial", "full"):
+ row_bufs["eval"].append(
+ {
+ **meta,
+ "ensembling": mode,
+ **evaluate_preds(val_gt, val_credences, mode).to_dict(),
+ **prompt_index_dict,
+ }
+ )
+
+ lr_dir = experiment_dir / "lr_models"
+ if not self.skip_supervised and lr_dir.exists():
+ with open(lr_dir / f"layer_{layer}.pt", "rb") as f:
+ lr_models = torch.load(f, map_location=device)
+ if not isinstance(
+ lr_models, list
+ ): # backward compatibility
+ lr_models = [lr_models]
+
+ for i, model in enumerate(lr_models):
+ model.eval()
+ row_bufs["lr_eval"].append(
+ {
+ "ensembling": mode,
+ "inlp_iter": i,
+ **meta,
+ **evaluate_preds(
+ val_gt, model(val_h), mode
+ ).to_dict(),
+ }
+ )
+
+ if isinstance(reporter, MultiReporter):
+ for i, res in enumerate(reporter.reporter_w_infos):
+ eval_all(res.model, res.prompt_index, i)
+ eval_all(reporter, "multi")
+ else:
+ eval_all(reporter)
 
  return {k: pd.DataFrame(v) for k, v in row_bufs.items()}

elk/run.py

@@ -30,6 +30,8 @@
  select_usable_devices,
 )
 
+PreparedData = dict[str, tuple[Tensor, Tensor, Tensor | None]]
+
 
 @dataclass
 class Run(ABC, Serializable):
@@ -46,11 +48,14 @@ class Run(ABC, Serializable):
  prompt_indices: tuple[int, ...] = ()
  """The indices of the prompt templates to use. If empty, all prompts are used."""
 
+ probe_per_prompt: bool = False
+ """If true, a probe is trained per prompt template. Otherwise, a single probe is
+ trained for all prompt templates."""
+
  concatenated_layer_offset: int = 0
  debug: bool = False
  min_gpu_mem: int | None = None # in bytes
  num_gpus: int = -1
- out_dir: Path | None = None
  disable_cache: bool = field(default=False, to_dict=False)
 
  def execute(
@@ -99,13 +104,16 @@ def execute(
  devices = select_usable_devices(self.num_gpus, min_memory=self.min_gpu_mem)
  num_devices = len(devices)
  func: Callable[[int], dict[str, pd.DataFrame]] = partial(
- self.apply_to_layer, devices=devices, world_size=num_devices
+ self.apply_to_layer,
+ devices=devices,
+ world_size=num_devices,
+ probe_per_prompt=self.probe_per_prompt,
  )
  self.apply_to_layers(func=func, num_devices=num_devices)
 
  @abstractmethod
  def apply_to_layer(
- self, layer: int, devices: list[str], world_size: int
+ self, layer: int, devices: list[str], world_size: int, probe_per_prompt: bool
  ) -> dict[str, pd.DataFrame]:
  """Train or eval a reporter on a single layer."""
 
@@ -125,7 +133,7 @@ def get_device(self, devices, world_size: int) -> str:
 
  def prepare_data(
  self, device: str, layer: int, split_type: Literal["train", "val"]
- ) -> dict[str, tuple[Tensor, Tensor, Tensor | None]]:
+ ) -> PreparedData:
  """Prepare data for the specified layer and split type."""
  out = {}
 
@@ -136,7 +144,7 @@ def prepare_data(
  labels = assert_type(Tensor, split["label"])
  hiddens = int16_to_float32(assert_type(Tensor, split[f"hidden_{layer}"]))
  if self.prompt_indices:
- hiddens = hiddens[:, self.prompt_indices]
+ hiddens = hiddens[:, self.prompt_indices, ...]
 
  with split.formatted_as("torch", device=device):
  has_preds = "model_logits" in split.features
@@ -186,7 +194,18 @@ def apply_to_layers(
  finally:
  # Make sure the CSVs are written even if we crash or get interrupted
  for name, dfs in df_buffers.items():
- df = pd.concat(dfs).sort_values(by=["layer", "ensembling"])
- df.round(4).to_csv(self.out_dir / f"{name}.csv", index=False)
+ sortby = ["layer", "ensembling"]
+ if "prompt_index" in dfs[0].columns:
+ sortby.append("prompt_index")
+ df = pd.concat(dfs).sort_values(by=sortby)
+
+ if "prompt_index" in df.columns:
+ cols = list(df.columns)
+ cols.insert(2, cols.pop(cols.index("prompt_index")))
+ df = df.reindex(columns=cols)
+
+ # Save the CSV
+ out_path = self.out_dir / f"{name}.csv"
+ df.round(4).to_csv(out_path, index=False)
  if self.debug:
  save_debug_log(self.datasets, self.out_dir)

elk/training/multi_reporter.py

@@ -0,0 +1,56 @@
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch as t
+
+from elk.training import CcsReporter
+from elk.training.common import Reporter
+
+SingleReporter = CcsReporter | Reporter
+
+
+@dataclass
+class ReporterWithInfo: # I don't love this name but I have no choice because
+ # of the other Reporter
+ model: SingleReporter
+ train_loss: float | None = None
+ prompt_index: int | None = None
+
+
+class MultiReporter:
+ def __init__(self, reporter: list[ReporterWithInfo]):
+ assert len(reporter) > 0, "Must have at least one reporter"
+ self.reporter_w_infos: list[ReporterWithInfo] = reporter
+ self.models = [r.model for r in reporter]
+ train_losses = (
+ [r.train_loss for r in reporter]
+ if reporter[0].train_loss is not None
+ else None
+ )
+
+ self.train_loss = (
+ sum(train_losses) / len(train_losses) # type: ignore
+ if train_losses is not None
+ else None
+ )
+
+ def __call__(self, h):
+ num_variants = h.shape[1]
+ assert len(self.models) == num_variants
+ credences = []
+ for i, reporter in enumerate(self.models):
+ credences.append(reporter(h[:, [i], :, :]))
+ return t.stack(credences, dim=0).mean(dim=0)
+
+ @staticmethod
+ def load(path: Path, layer: int, device: str):
+ prompt_folders = [p for p in path.iterdir() if p.is_dir()]
+ reporters = [
+ (
+ t.load(folder / "reporters" / f"layer_{layer}.pt", map_location=device),
+ int(folder.name.split("_")[-1]), # prompt index
+ )
+ for folder in prompt_folders
+ ]
+ # we don't care about the train losses for evaluating
+ return MultiReporter([ReporterWithInfo(r, None, pi) for r, pi in reporters])