Support --binarize again

elk/evaluation/evaluate.py

@@ -6,7 +6,7 @@
 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
@@ -25,16 +25,13 @@ class Eval(Run):
  def __post_init__(self):
  assert self.source, "Must specify a source experiment."
 
- # 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
- )
+ if not self.out_dir:
+ self.out_dir = self.source / "transfer" / "+".join(self.data.datasets)
 
  def execute(self, highlight_color: str = "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/extraction.py

@@ -27,6 +27,7 @@
 
 from ..promptsource import DatasetTemplates
 from ..utils import (
+ Color,
  assert_type,
  colorize,
  float32_to_int16,
@@ -60,6 +61,9 @@ class Extract(Serializable):
  data_dirs: tuple[str, ...] = ()
  """Directory to use for caching the hiddens. Defaults to `HF_DATASETS_CACHE`."""
 
+ binarize: bool = False
+ """Whether to binarize the dataset labels for multi-class datasets."""
+
  max_examples: tuple[int, int] = (1000, 1000)
  """Maximum number of examples to use from each split of the dataset."""
 
@@ -70,16 +74,16 @@ class Extract(Serializable):
  """The number of prompt templates to use for each example. If -1, all available
  templates are used."""
 
- seed: int = 42
- """Seed to use for prompt randomization. Defaults to 42."""
-
  layers: tuple[int, ...] = ()
  """Indices of layers to extract hidden states from. We follow the HF convention, so
  0 is the embedding, and 1 is the output of the first transformer layer."""
 
  layer_stride: InitVar[int] = 1
  """Shortcut for `layers = (0,) + tuple(range(1, num_layers + 1, stride))`."""
 
+ seed: int = 42
+ """Seed to use for prompt randomization. Defaults to 42."""
+
  template_path: str | None = None
  """Path to pass into `DatasetTemplates`. By default we use the dataset name."""
 
@@ -170,6 +174,7 @@ def extract_hiddens(
 
  prompt_ds = load_prompts(
  ds_names[0],
+ binarize=cfg.binarize,
  split_type=split_type,
  template_path=cfg.template_path,
  rank=rank,
@@ -322,8 +327,10 @@ def hidden_features(cfg: Extract) -> tuple[DatasetInfo, Features]:
 
  ds_features = assert_type(Features, info.features)
  label_col = prompter.label_column or infer_label_column(ds_features)
- num_classes = len(prompter.label_choices) or infer_num_classes(
- ds_features[label_col]
+ num_classes = (
+ 2
+ if cfg.binarize
+ else (len(prompter.label_choices) or infer_num_classes(ds_features[label_col]))
  )
 
  num_variants = cfg.num_variants
@@ -369,7 +376,7 @@ def extract(
  cfg: "Extract",
  *,
  disable_cache: bool = False,
- highlight_color: str = "cyan",
+ highlight_color: Color = "cyan",
  num_gpus: int = -1,
  min_gpu_mem: int | None = None,
  split_type: Literal["train", "val", None] = None,

elk/extraction/prompt_loading.py

@@ -16,6 +16,7 @@
 def load_prompts(
  ds_string: str,
  *,
+ binarize: bool = False,
  num_shots: int = 0,
  num_variants: int = -1,
  seed: int = 42,
@@ -28,6 +29,7 @@ def load_prompts(
 
  Args:
  ds_string: Name of HF dataset to use, e.g. `"super_glue:boolq"` or `"imdb"`.
+ binarize: Whether to binarize the dataset labels for multi-class datasets.
  num_shots: The number of examples to use in few-shot prompts. If zero, prompts
  are zero-shot.
  seed: The seed to use for prompt randomization.
@@ -108,6 +110,7 @@ def load_prompts(
  for example in ds:
  yield _convert_to_prompts(
  example,
+ binarize=binarize,
  label_column=label_column,
  label_choices=label_choices, # type: ignore[arg-type]
  num_variants=num_variants,
@@ -120,6 +123,7 @@ def load_prompts(
 def _convert_to_prompts(
  example: dict[str, Any],
  prompter: DatasetTemplates,
+ binarize: bool,
  label_column: str,
  label_choices: list[bool | int | str],
  num_variants: int,
@@ -140,6 +144,15 @@ def qa_cat(q: str, a: str) -> str:
  # For sanity checking that prompts are unique
  prompt_counter = Counter()
  label = example[label_column]
+ if binarize:
+ # Replace the full list of possibilities with a randomly sampled false label
+ # and the correct label, as done in the DLK paper. Note that this does add some
+ # "supervision" by stacking the deck in favor of the correct answer.
+ label_choices = [
+ rng.choice([c for c in label_choices if c != label]),
+ label,
+ ]
+ rng.shuffle(label_choices)
 
  for template in templates:
  choices = []

elk/files.py

@@ -38,8 +38,3 @@ def memorably_named_dir(parent: Path):
  out_dir = parent / sub_dir
  out_dir.mkdir(parents=True, exist_ok=True)
  return out_dir
-
-
-def transfer_eval_directory(source: Path) -> Path:
- """Return the directory where transfer evals are stored."""
- return elk_reporter_dir() / source / "transfer_eval"