merge

elk/__init__.py

@@ -1 +1 @@
-from .extraction import extract_hiddens, ExtractionConfig
+from .extraction import extract_hiddens, Extract

elk/__main__.py

@@ -1,67 +1,32 @@
 """Main entry point for `elk`."""
 
-from .extraction import extract, ExtractionConfig
-from .evaluation.evaluate import EvaluateConfig, evaluate_reporters
-from .training import RunConfig
-from .training.train import train
+from dataclasses import dataclass
 from pathlib import Path
+from typing import Union
+
 from simple_parsing import ArgumentParser
 
+from elk.evaluation.evaluate import Eval
+from elk.extraction.extraction import Extract
+from elk.training.train import Elicit
 
-def run():
- parser = ArgumentParser(add_help=False)
- subparsers = parser.add_subparsers(dest="command", required=True)
 
- extract_parser = subparsers.add_parser(
- "extract", help="Extract hidden states from a model."
- )
- extract_parser.add_arguments(ExtractionConfig, dest="extraction")
- extract_parser.add_argument(
- "--output",
- "-o",
- type=Path,
- help="Path to save hidden states to.",
- required=True,
- )
- extract_parser.add_argument(
- "--num_gpus",
- type=int,
- help="Maximum number of GPUs to use.",
- required=False,
- default=-1,
- )
+@dataclass
+class Command:
+ """Some top-level command"""
 
- elicit_parser = subparsers.add_parser(
- "elicit",
- help=(
- "Extract and train a set of ELK reporters "
- "on hidden states from `elk extract`. "
- ),
- conflict_handler="resolve",
- )
- elicit_parser.add_arguments(RunConfig, dest="run")
- elicit_parser.add_argument(
- "--output",
- "-o",
- type=Path,
- help="Path to save checkpoints to.",
- )
+ command: Union[Elicit, Eval, Extract]
 
- subparsers.add_parser(
- "eval",
- help="Evaluate a set of ELK reporters generated by `elk train`.",
- ).add_arguments(EvaluateConfig, dest="eval")
+ def execute(self):
+ return self.command.execute()
 
- args = parser.parse_args()
 
- if args.command == "extract":
- extract(args.extraction, args.num_gpus).save_to_disk(args.output)
- elif args.command == "elicit":
- train(args.run, args.output)
- elif args.command == "eval":
- evaluate_reporters(args.eval)
- else:
- raise ValueError(f"Unknown command {args.command}")
+def run():
+ parser = ArgumentParser(add_help=False)
+ parser.add_arguments(Command, dest="run")
+ args = parser.parse_args()
+ run: Command = args.run
+ run.execute()
 
 
 if __name__ == "__main__":

elk/evaluation/evaluate.py

@@ -6,13 +6,13 @@
 from simple_parsing.helpers import Serializable, field
 from torch import Tensor
 from tqdm.auto import tqdm
-from typing import Literal, Optional, cast
+from typing import Literal, Optional, cast, Callable
 import csv
 import os
 import torch
 import torch.multiprocessing as mp
 
-from ..extraction import ExtractionConfig, extract
+from elk.extraction.extraction import Extract
 from ..files import elk_reporter_dir, memorably_named_dir
 from ..utils import (
  assert_type,
@@ -21,96 +21,91 @@
  select_usable_devices,
 )
 
+import torch
+from simple_parsing import Serializable, field
+
+from elk.files import elk_reporter_dir
+from elk.run import Run
+from elk.training import Reporter
+from elk.evaluation.evaluate_log import EvalLog
+from elk.utils import select_usable_devices
+
 
 @dataclass
-class EvaluateConfig(Serializable):
- target: ExtractionConfig
+class Eval(Serializable):
+ """
+ Full specification of a reporter evaluation run.
+
+ Args:
+ data: Config specifying hidden states on which the reporter will be evaluated.
+ source: The name of the source run directory
+ which contains the reporters directory.
+ normalization: The normalization method to use. Defaults to "meanonly". See
+ `elk.training.preprocessing.normalize()` for details.
+ num_gpus: The number of GPUs to use. Defaults to -1, which means
+ "use all available GPUs".
+ debug: When in debug mode, a useful log file is saved to the memorably-named
+ output directory. Defaults to False.
+ """
+
+ data: Extract
  source: str = field(positional=True)
  normalization: Literal["legacy", "none", "elementwise", "meanonly"] = "meanonly"
- num_gpus: int = -1
-
-
-def evaluate_reporter(
- cfg: EvaluateConfig,
- dataset: DatasetDict,
- layer: int,
- devices: list[str],
- world_size: int = 1,
-):
- """Evaluate a single reporter on a single layer."""
- rank = os.getpid() % world_size
- device = devices[rank]
-
- # Note: currently we're just upcasting to float32 so we don't have to deal with
- # grad scaling (which isn't supported for LBFGS), while the hidden states are
- # saved in float16 to save disk space. In the future we could try to use mixed
- # precision training in at least some cases.
- with dataset.formatted_as("torch", device=device, dtype=torch.int16):
- train_split, val_split = select_train_val_splits(dataset)
- train, val = dataset[train_split], dataset[val_split]
- test_labels = cast(Tensor, val["label"])
-
- _, test_h = normalize(
- int16_to_float32(assert_type(Tensor, train[f"hidden_{layer}"])),
- int16_to_float32(assert_type(Tensor, val[f"hidden_{layer}"])),
- method=cfg.normalization,
- )
-
- reporter_path = elk_reporter_dir() / cfg.source / "reporters" / f"layer_{layer}.pt"
- reporter = torch.load(reporter_path, map_location=device)
- reporter.eval()
-
- test_x0, test_x1 = test_h.unbind(dim=-2)
 
- test_result = reporter.score(test_labels, test_x0, test_x1)
-
- stats = [layer, *test_result]
- return stats
+ debug: bool = False
+ out_dir: Optional[Path] = None
+ num_gpus: int = -1
 
+ def execute(self):
+ transfer_eval = elk_reporter_dir() / self.source / "transfer_eval"
 
-def evaluate_reporters(cfg: EvaluateConfig, out_dir: Optional[Path] = None):
- ds = extract(cfg.target, num_gpus=cfg.num_gpus)
+ run = Evaluate(cfg=self, out_dir=transfer_eval)
+  run.evaluate()
 
- layers = [
- int(feat[len("hidden_") :])
- for feat in ds["train"].features
- if feat.startswith("hidden_")
- ]
 
- devices = select_usable_devices(cfg.num_gpus, min_memory=cfg.target.min_gpu_mem)
- num_devices = len(devices)
+@dataclass
+class Evaluate(Run):
+ cfg: Eval
+
+ def evaluate_reporter(
+ self, layer: int, devices: list[str], world_size: int = 1
+ ) -> EvalLog:
+ """Evaluate a single reporter on a single layer."""
+ device = self.get_device(devices, world_size)
+
+ _, _, test_x0, test_x1, _, test_labels = self.prepare_data(
+ device,
+ layer,
+ )
 
- transfer_eval = elk_reporter_dir() / cfg.source / "transfer_eval"
- transfer_eval.mkdir(parents=True, exist_ok=True)
+ reporter_path = (
+ elk_reporter_dir() / self.cfg.source / "reporters" / f"layer_{layer}.pt"
+ )
+ reporter: Reporter = torch.load(reporter_path, map_location=device)
+ reporter.eval()
 
- if out_dir is None:
- out_dir = memorably_named_dir(transfer_eval)
- else:
- out_dir.mkdir(parents=True, exist_ok=True)
+ test_result = reporter.score(
+ test_labels,
+ test_x0,
+ test_x1,
+ )
 
- # Print the output directory in bold with escape codes
- print(f"Saving results to \033[1m{out_dir}\033[0m")
+ return EvalLog(
+ layer=layer,
+ eval_result=test_result,
+ )
 
- with open(out_dir / "cfg.yaml", "w") as f:
- cfg.dump_yaml(f)
+ def evaluate(self):
+ """Evaluate the reporter on all layers."""
+ devices = select_usable_devices(self.cfg.num_gpus, min_memory=cfg.target.min_gpu_mem)
 
- cols = ["layer", "loss", "acc", "cal_acc", "auroc"]
- # Evaluate reporters for each layer in parallel
- with mp.Pool(num_devices) as pool, open(out_dir / "eval.csv", "w") as f:
- fn = partial(
- evaluate_reporter, cfg, ds, devices=devices, world_size=num_devices
+ num_devices = len(devices)
+ func: Callable[[int], EvalLog] = partial(
+ self.evaluate_reporter, devices=devices, world_size=num_devices
+ )
+ self.apply_to_layers(
+ func=func,
+ num_devices=num_devices,
+ to_csv_line=lambda item: item.to_csv_line(),
+ csv_columns=EvalLog.csv_columns(),
  )
- writer = csv.writer(f)
- writer.writerow(cols)
-
- mapper = pool.imap_unordered if num_devices > 1 else map
- row_buf = []
- try:
- for i, *stats in tqdm(mapper(fn, layers), total=len(layers)):
- row_buf.append([i] + [f"{s:.4f}" for s in stats])
- finally:
- # Make sure the CSV is written even if we crash or get interrupted
- for row in sorted(row_buf):
- writer.writerow(row)
-
- print("Results saved")

elk/evaluation/evaluate_log.py

@@ -0,0 +1,27 @@
+from dataclasses import dataclass
+
+from elk.training.reporter import EvalResult
+
+
+@dataclass
+class EvalLog:
+ """The result of running eval on a layer of a dataset"""
+
+ layer: int
+ eval_result: EvalResult
+
+ @staticmethod
+ def csv_columns() -> list[str]:
+ return ["layer", "acc", "cal_acc", "auroc", "ece"]
+
+ def to_csv_line(self) -> list[str]:
+ items = [
+ self.layer,
+ self.eval_result.acc,
+ self.eval_result.cal_acc,
+ self.eval_result.auroc,
+ self.eval_result.ece,
+ ]
+ return [
+ f"{item:.4f}" if isinstance(item, float) else str(item) for item in items
+ ]

elk/extraction/__init__.py

@@ -1,4 +1,4 @@
 from .balanced_sampler import BalancedSampler, FewShotSampler
-from .extraction import ExtractionConfig, extract_hiddens, extract
+from .extraction import Extract, extract_hiddens, extract
 from .generator import _GeneratorConfig, _GeneratorBuilder
-from .prompt_loading import PromptConfig, load_prompts
+from .prompt_loading import PromptConfig, load_prompts