Binarized meta-templates; some extraction refactoring

elk/evaluation/evaluate.py

@@ -1,33 +1,33 @@
 from collections import defaultdict
 from dataclasses import dataclass
+from pathlib import Path
 
 import pandas as pd
 import torch
 from simple_parsing.helpers import field
 
-from ..files import elk_reporter_dir, transfer_eval_directory
+from ..files import elk_reporter_dir
 from ..metrics import evaluate_preds
 from ..run import Run
 from ..training import Reporter
+from ..utils import Color
 
 
-@dataclass
+@dataclass(kw_only=True)
 class Eval(Run):
  """Full specification of a reporter evaluation run."""
 
- source: str = field(default="", positional=True)
+ source: Path = field(positional=True)
  skip_supervised: bool = False
 
  def __post_init__(self):
- assert self.source, "Must specify a source experiment."
+ if not self.out_dir:
+ self.out_dir = self.source / "transfer" / "+".join(self.data.datasets)
 
- # Set the output directory to the transfer directory if it's not specified
- self.out_dir = (
- transfer_eval_directory(self.source)
- if self.out_dir is None
- else self.out_dir
- )
+ def execute(self, highlight_color: Color = "cyan"):
+ return super().execute(highlight_color, split_type="val")
 
+ @torch.inference_mode()
  def apply_to_layer(
  self, layer: int, devices: list[str], world_size: int
  ) -> dict[str, pd.DataFrame]:

elk/extraction/__init__.py

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

elk/extraction/balanced_sampler.py

@@ -1,8 +1,8 @@
 from collections import deque
-from dataclasses import dataclass, field
+from dataclasses import InitVar, dataclass, field
 from itertools import cycle
 from random import Random
-from typing import Iterable, Iterator, Optional
+from typing import Hashable, Iterable, Iterator, Optional
 
 from datasets import Features, IterableDataset
 from torch.utils.data import IterableDataset as TorchIterableDataset
@@ -26,25 +26,29 @@ class BalancedSampler(TorchIterableDataset):
  """
 
  data: Iterable[dict]
- num_classes: int
+ label_choices: InitVar[set[Hashable]]
  buffer_size: int = 1000
- buffers: dict[int, deque[dict]] = field(default_factory=dict, init=False)
+ buffers: dict[Hashable, deque[dict]] = field(default_factory=dict, init=False)
  label_col: str = "label"
+ strict: bool = True
 
- def __post_init__(self):
+ def __post_init__(self, label_choices: set[Hashable]):
  # Initialize empty buffers
  self.buffers = {
- label: deque(maxlen=self.buffer_size) for label in range(self.num_classes)
+ label: deque(maxlen=self.buffer_size) for label in label_choices
  }
 
  def __iter__(self):
  for sample in self.data:
  label = sample[self.label_col]
-
- # This whole class is a no-op if the label is not an integer
- if not isinstance(label, int):
- yield sample
- continue
+ if label not in self.buffers:
+ if self.strict:
+ raise ValueError(
+ f"Expected label to be one of {self.buffers}, got {label}"
+ )
+ else:
+ # Just skip this sample
+ continue
 
  # Add the sample to the buffer for its class label
  self.buffers[label].append(sample)