Various fixes

elk/evaluation/evaluate.py

@@ -11,7 +11,7 @@
 from ..files import elk_reporter_dir
 from ..run import Run
 from ..training import Reporter
-from ..training.baseline import evaluate_baseline
+from ..training.supervised import evaluate_supervised
 from ..utils import select_usable_devices
 
 
@@ -61,11 +61,7 @@ def evaluate_reporter(
  ) -> 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_output = self.prepare_data(device, layer, "val")
 
  experiment_dir = elk_reporter_dir() / self.cfg.source
 
@@ -94,7 +90,7 @@ def evaluate_reporter(
  with open(lr_dir / f"layer_{layer}.pt", "rb") as f:
  lr_model = torch.load(f, map_location=device).eval()
 
- lr_auroc, lr_acc = evaluate_baseline(lr_model, val_x0, val_x1, val_gt)
+ lr_auroc, lr_acc = evaluate_supervised(lr_model, val_x0, val_x1, val_gt)
 
  stats_row["lr_auroc"] = lr_auroc
  stats_row["lr_acc"] = lr_acc

elk/extraction/extraction.py

@@ -23,6 +23,7 @@
 from transformers import AutoConfig, AutoTokenizer
 from transformers.modeling_outputs import Seq2SeqLMOutput
 
+from ..promptsource import DatasetTemplates
 from ..utils import (
  assert_type,
  convert_span,
@@ -257,8 +258,9 @@ def get_splits() -> SplitDict:
  available_splits = assert_type(SplitDict, info.splits)
  train_name, val_name = select_train_val_splits(available_splits)
  print(
- f"{info.builder_name}: using '{train_name}' for training and '{val_name}'"
- f" for validation"
+ # Cyan color for dataset name
+ f"\033[36m{info.builder_name}\033[0m: using '{train_name}' for training and"
+ f" '{val_name}' for validation"
  )
  limit_list = cfg.prompts.max_examples
 
@@ -275,11 +277,15 @@ def get_splits() -> SplitDict:
  )
 
  model_cfg = AutoConfig.from_pretrained(cfg.model)
- num_variants = cfg.prompts.num_variants
 
  ds_name, _, config_name = cfg.prompts.datasets[0].partition(" ")
  info = get_dataset_config_info(ds_name, config_name or None)
 
+ num_variants = cfg.prompts.num_variants
+ if num_variants < 0:
+ prompter = DatasetTemplates(ds_name, config_name)
+ num_variants = len(prompter.templates)
+
  layer_cols = {
  f"hidden_{layer}": Array3D(
  dtype="int16",

elk/run.py

@@ -3,14 +3,14 @@
 from abc import ABC
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import TYPE_CHECKING, Callable, Union
+from typing import TYPE_CHECKING, Callable, Literal, Union
 
 import numpy as np
 import pandas as pd
 import torch
 import torch.multiprocessing as mp
 import yaml
-from datasets import DatasetDict, concatenate_datasets
+from datasets import DatasetDict
 from torch import Tensor
 from tqdm import tqdm
 
@@ -84,40 +84,29 @@ def get_device(self, devices, world_size: int) -> str:
  device = devices[rank]
  return device
 
- def prepare_data(self, device: str, layer: int) -> tuple:
- """Prepare the data for training and validation."""
+ def prepare_data(
+ self, device: str, layer: int, split_type: Literal["train", "val"]
+ ) -> dict[str, tuple[Tensor, Tensor, Tensor, np.ndarray | None]]:
+ """Prepare data for the specified layer and split type."""
+ out = {}
 
- train_sets = []
- val_output = {}
-
- # We handle train and val differently. We want to concatenate all of the
- # train sets together, but we want to keep the val sets separate so that we can
- # compute evaluation metrics separately for each dataset.
  for ds in self.datasets:
- train_split, val_split = select_train_val_splits(ds)
- train_sets.append(ds[train_split])
+ train_name, val_name = select_train_val_splits(ds)
+ key = train_name if split_type == "train" else val_name
 
- val = ds[val_split].with_format("torch", device=device, dtype=torch.int16)
- val_labels = assert_type(Tensor, val["label"])
- val_h = int16_to_float32(assert_type(torch.Tensor, val[f"hidden_{layer}"]))
- val_x0, val_x1 = val_h.unbind(dim=-2)
+ split = ds[key].with_format("torch", device=device, dtype=torch.int16)
+ labels = assert_type(Tensor, split["label"])
+ val_h = int16_to_float32(assert_type(Tensor, split[f"hidden_{layer}"]))
+ x0, x1 = val_h.unbind(dim=-2)
 
- with val.formatted_as("numpy"):
- has_preds = "model_preds" in val.features
- val_lm_preds = val["model_preds"] if has_preds else None
+ with split.formatted_as("numpy"):
+ has_preds = "model_preds" in split.features
+ lm_preds = split["model_preds"] if has_preds else None
 
  ds_name = get_dataset_name(ds)
- val_output[ds_name] = (val_x0, val_x1, val_labels, val_lm_preds)
-
- train = concatenate_datasets(train_sets).with_format(
- "torch", device=device, dtype=torch.int16
- )
-
- train_labels = assert_type(Tensor, train["label"])
- train_h = int16_to_float32(assert_type(torch.Tensor, train[f"hidden_{layer}"]))
- x0, x1 = train_h.unbind(dim=-2)
+ out[ds_name] = (x0, x1, labels, lm_preds)
 
- return x0, x1, train_labels, val_output
+ return out
 
  def concatenate(self, layers):
  """Concatenate hidden states from a previous layer."""

elk/training/eigen_reporter.py

@@ -2,12 +2,11 @@
 
 from dataclasses import dataclass
 from typing import Optional
-from warnings import warn
 
 import torch
 from torch import Tensor, nn, optim
 
-from ..truncated_eigh import ConvergenceError, truncated_eigh
+from ..truncated_eigh import truncated_eigh
 from ..utils.math_util import cov_mean_fused
 from .reporter import Reporter, ReporterConfig
 
@@ -184,22 +183,17 @@ def update(self, x_pos: Tensor, x_neg: Tensor) -> None:
  self.contrastive_xcov_M2.addmm_(neg_delta.mT, pos_delta2)
  self.contrastive_xcov_M2.addmm_(pos_delta.mT, neg_delta2)
 
- def fit_streaming(self) -> float:
+ def fit_streaming(self, truncated: bool = False) -> float:
  """Fit the probe using the current streaming statistics."""
  A = (
  self.config.var_weight * self.intercluster_cov
  - self.config.inv_weight * self.intracluster_cov
  - self.config.neg_cov_weight * self.contrastive_xcov
  )
 
- try:
+ if truncated:
  L, Q = truncated_eigh(A, k=self.config.num_heads)
- except (ConvergenceError, RuntimeError):
- warn(
- "Truncated eigendecomposition failed to converge. Falling back on "
- "PyTorch's dense eigensolver."
- )
-
+ else:
  L, Q = torch.linalg.eigh(A)
  L, Q = L[-self.config.num_heads :], Q[:, -self.config.num_heads :]
 

elk/training/supervised.py

@@ -0,0 +1,47 @@
+import torch
+from einops import rearrange, repeat
+from sklearn.metrics import accuracy_score, roc_auc_score
+from torch import Tensor
+
+from ..utils import assert_type
+from .classifier import Classifier
+
+
+def evaluate_supervised(
+ lr_model: Classifier, val_x0: Tensor, val_x1: Tensor, val_labels: Tensor
+) -> tuple[float, float]:
+ X = torch.cat([val_x0, val_x1])
+ d = X.shape[-1]
+ X_val = X.view(-1, d)
+ with torch.no_grad():
+ lr_preds = lr_model(X_val).sigmoid().cpu()
+
+ val_labels_aug = (
+ torch.cat([val_labels, 1 - val_labels]).repeat_interleave(val_x0.shape[1])
+ ).cpu()
+
+ lr_acc = accuracy_score(val_labels_aug, lr_preds > 0.5)
+ lr_auroc = roc_auc_score(val_labels_aug, lr_preds)
+
+ return assert_type(float, lr_auroc), assert_type(float, lr_acc)
+
+
+def train_supervised(data: dict[str, tuple], device: str) -> Classifier:
+ Xs, train_labels = [], []
+
+ for x0, x1, labels, _ in data.values():
+ (_, v, _) = x0.shape
+ x0 = rearrange(x0, "n v d -> (n v) d")
+ x1 = rearrange(x1, "n v d -> (n v) d")
+
+ labels = repeat(labels, "n -> (n v)", v=v)
+ labels = torch.cat([labels, 1 - labels])
+
+ Xs.append(torch.cat([x0, x1]).squeeze())
+ train_labels.append(labels)
+
+ X, train_labels = torch.cat(Xs), torch.cat(train_labels)
+ lr_model = Classifier(X.shape[-1], device=device)
+ lr_model.fit_cv(X, train_labels)
+
+ return lr_model