isort

elk/__init__.py

@@ -1,3 +1,3 @@
-from .extraction import extract_hiddens, Extract
+from .extraction import Extract, extract_hiddens
 
 __all__ = ["extract_hiddens", "Extract"]

elk/calibration.py

@@ -1,8 +1,9 @@
+import warnings
 from dataclasses import dataclass, field
-from torch import Tensor
 from typing import NamedTuple
+
 import torch
-import warnings
+from torch import Tensor
 
 
 class CalibrationEstimate(NamedTuple):

elk/eigsh.py

@@ -1,7 +1,8 @@
-from torch import Tensor
 from typing import Literal, Optional
+
 import torch
 import torch.nn.functional as F
+from torch import Tensor
 
 
 def lanczos_eigsh(

elk/evaluation/evaluate.py

@@ -1,15 +1,17 @@
-from ..extraction import Extract
-from ..files import elk_reporter_dir
-from ..run import Run
-from ..training import Reporter
-from ..utils import select_usable_devices
-from .evaluate_log import EvalLog
 from dataclasses import dataclass
 from functools import partial
 from pathlib import Path
-from simple_parsing.helpers import Serializable, field
 from typing import Callable, Literal, Optional
+
+import pandas as pd
 import torch
+from simple_parsing.helpers import Serializable, field
+
+from ..extraction import Extract
+from ..files import elk_reporter_dir
+from ..run import Run
+from ..training import Reporter
+from ..utils import select_usable_devices
 
 
 @dataclass
@@ -52,7 +54,7 @@ class Evaluate(Run):
 
  def evaluate_reporter(
  self, layer: int, devices: list[str], world_size: int = 1
- ) -> EvalLog:
+ ) -> pd.Series:
  """Evaluate a single reporter on a single layer."""
  device = self.get_device(devices, world_size)
 
@@ -73,9 +75,11 @@ def evaluate_reporter(
  test_x1,
  )
 
- return EvalLog(
- layer=layer,
- eval_result=test_result,
+ return pd.Series(
+ {
+ "layer": layer,
+ **test_result._asdict(),
+ }
  )
 
  def evaluate(self):
@@ -85,12 +89,10 @@ def evaluate(self):
  )
 
  num_devices = len(devices)
- func: Callable[[int], EvalLog] = partial(
+ func: Callable[[int], pd.Series] = 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(),
  )

elk/extraction/__init__.py

@@ -1,6 +1,6 @@
 from .balanced_sampler import BalancedSampler, FewShotSampler
-from .extraction import Extract, extract_hiddens, extract
-from .generator import _GeneratorConfig, _GeneratorBuilder
+from .extraction import Extract, extract, extract_hiddens
+from .generator import _GeneratorBuilder, _GeneratorConfig
 from .prompt_loading import PromptConfig, load_prompts
 
 __all__ = [

elk/extraction/balanced_sampler.py

@@ -1,12 +1,14 @@
-from ..math_util import stochastic_round_constrained
-from ..utils import infer_label_column
-from ..utils.typing import assert_type
 from collections import deque
-from datasets import IterableDataset, Features
 from itertools import cycle
 from random import Random
+from typing import Iterable, Iterator, Optional
+
+from datasets import Features, IterableDataset
 from torch.utils.data import IterableDataset as TorchIterableDataset
-from typing import Iterator, Optional, Iterable
+
+from ..math_util import stochastic_round_constrained
+from ..utils import infer_label_column
+from ..utils.typing import assert_type
 
 
 class BalancedSampler(TorchIterableDataset):

elk/extraction/extraction.py

@@ -1,5 +1,11 @@
 """Functions for extracting the hidden states of a model."""
+import logging
+import os
 from dataclasses import InitVar, dataclass
+from itertools import islice
+from typing import Iterable, Literal, Optional, Union
+
+import torch
 from datasets import (
  Array3D,
  ClassLabel,
@@ -11,18 +17,12 @@
  Value,
  get_dataset_config_info,
 )
-from itertools import islice
 from simple_parsing import Serializable, field
 from torch import Tensor
 from transformers import AutoConfig, AutoTokenizer
 from transformers.modeling_outputs import Seq2SeqLMOutput
-from typing import Iterable, Literal, Optional, Union
-import logging
-import os
-import torch
 
 # import torch.nn.functional as F
-
 from ..utils import (
  assert_type,
  convert_span,

elk/extraction/generator.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field
-from typing import Callable, Optional, Any, Dict
+from typing import Any, Callable, Dict, Optional
 
 import datasets
 from datasets.splits import NamedSplit

elk/extraction/prompt_loading.py

@@ -1,3 +1,15 @@
+from dataclasses import dataclass
+from random import Random
+from typing import Any, Iterator, Literal, Optional
+
+from datasets import (
+ Dataset,
+ Features,
+ load_dataset,
+)
+from datasets.distributed import split_dataset_by_node
+from simple_parsing.helpers import Serializable, field
+
 from ..promptsource import DatasetTemplates
 from ..utils import (
  assert_type,
@@ -7,16 +19,6 @@
  select_train_val_splits,
 )
 from .balanced_sampler import FewShotSampler
-from dataclasses import dataclass
-from datasets import (
- load_dataset,
- Dataset,
- Features,
-)
-from datasets.distributed import split_dataset_by_node
-from random import Random
-from simple_parsing.helpers import field, Serializable
-from typing import Any, Iterator, Literal, Optional
 
 
 @dataclass

elk/files.py

@@ -1,13 +1,13 @@
 """Helper functions for dealing with files."""
 
-from pathlib import Path
 import json
 import os
 import random
+from pathlib import Path
 from typing import Optional
 
-from simple_parsing import Serializable
 import yaml
+from simple_parsing import Serializable
 
 
 def elk_reporter_dir() -> Path:

elk/logging.py

@@ -1,4 +1,5 @@
 import logging
+
 from .utils import select_train_val_splits
 
 

elk/math_util.py

@@ -1,7 +1,8 @@
-from torch import Tensor
 import math
 import random
+
 import torch
+from torch import Tensor
 
 
 @torch.jit.script

elk/parsing.py

@@ -1,4 +1,5 @@
 import re
+
 from .training.losses import LOSSES
 
 

elk/promptsource/templates.py

@@ -1,14 +1,14 @@
-from collections import Counter, defaultdict
-from jinja2 import BaseLoader, Environment, meta
-from pathlib import Path
-from shutil import rmtree
-from typing import Optional
 import logging
 import os
 import random
 import uuid
-import yaml
+from collections import Counter, defaultdict
+from pathlib import Path
+from shutil import rmtree
+from typing import Optional
 
+import yaml
+from jinja2 import BaseLoader, Environment, meta
 
 # Truncation of jinja template variables
 # 1710 = 300 words x 4.7 avg characters per word + 300 spaces

elk/run.py

@@ -5,29 +5,29 @@
 from pathlib import Path
 from typing import (
  TYPE_CHECKING,
+ Callable,
  Optional,
  Union,
- Callable,
- Iterator,
 )
 
 import numpy as np
+import pandas as pd
 import torch
 import torch.multiprocessing as mp
 from datasets import DatasetDict
 from torch import Tensor
 from tqdm import tqdm
 
-from elk.extraction.extraction import extract
-from elk.files import create_output_directory, save_config, save_meta
-from elk.training.preprocessing import normalize
-from elk.utils.csv import write_iterator_to_file, Log
-from elk.utils.data_utils import get_layers, select_train_val_splits
-from elk.utils.typing import assert_type, int16_to_float32
+from .extraction import extract
+from .files import create_output_directory, save_config, save_meta
+from .logging import save_debug_log
+from .training.preprocessing import normalize
+from .utils import assert_type, int16_to_float32
+from .utils.data_utils import get_layers, select_train_val_splits
 
 if TYPE_CHECKING:
- from elk.evaluation.evaluate import Eval
- from elk.training.train import Elicit
+ from .evaluation.evaluate import Eval
+ from .training.train import Elicit
 
 
 @dataclass
@@ -103,10 +103,8 @@ def concatenate(self, layers):
 
  def apply_to_layers(
  self,
- func: Callable[[int], Log],
+ func: Callable[[int], pd.Series],
  num_devices: int,
- to_csv_line: Callable[[Log], list[str]],
- csv_columns: list[str],
  ):
  """Apply a function to each layer of the dataset in parallel
  and writes the results to a CSV file.
@@ -115,10 +113,7 @@ def apply_to_layers(
  func: The function to apply to each layer.
  The int is the index of the layer.
  num_devices: The number of devices to use.
- to_csv_line: A function that converts a Log to a list of strings.
- This has to be injected in because the Run class does not know
- the extra options e.g. skip_baseline to apply to function.
- csv_columns: The columns of the CSV file."""
+ """
  self.out_dir = assert_type(Path, self.out_dir)
 
  layers: list[int] = get_layers(self.dataset)
@@ -129,15 +124,14 @@ def apply_to_layers(
  # Should we write to different CSV files for elicit vs eval?
  with mp.Pool(num_devices) as pool, open(self.out_dir / "eval.csv", "w") as f:
  mapper = pool.imap_unordered if num_devices > 1 else map
- iterator: Iterator[Log] = tqdm( # type: ignore
- mapper(func, layers), total=len(layers)
- )
- write_iterator_to_file(
- iterator=iterator,
- file=f,
- debug=self.cfg.debug,
- dataset=self.dataset,
- out_dir=self.out_dir,
- csv_columns=csv_columns,
- to_csv_line=to_csv_line,
- )
+ row_buf = []
+
+ try:
+ for row in tqdm(mapper(func, layers), total=len(layers)):
+ row_buf.append(row)
+ finally:
+ # Make sure the CSV is written even if we crash or get interrupted
+ df = pd.DataFrame(row_buf).sort_values(by="layer")
+ df.to_csv(f, index=False)
+ if self.cfg.debug:
+ save_debug_log(self.dataset, self.out_dir)

elk/training/__init__.py

@@ -2,7 +2,6 @@
 from .eigen_reporter import EigenReporter, EigenReporterConfig
 from .reporter import OptimConfig, Reporter, ReporterConfig
 
-
 __all__ = [
  "Reporter",
  "ReporterConfig",