Fix bugs in binary case

elk/extraction/extraction.py

@@ -33,7 +33,6 @@
  select_train_val_splits,
  select_usable_devices,
 )
-from .balanced_sampler import BalancedSampler
 from .generator import _GeneratorBuilder
 from .prompt_loading import PromptConfig, load_prompts
 
@@ -126,7 +125,7 @@ def extract_hiddens(
  if rank == world_size - 1:
  max_examples += global_max_examples % world_size
 
- for example in islice(BalancedSampler(prompt_ds, 3), max_examples):
+ for example in islice(prompt_ds, max_examples):
  num_variants = len(example["prompts"])
  num_choices = len(example["prompts"][0])
 

elk/extraction/prompt_loading.py

@@ -1,11 +1,11 @@
 from collections import Counter
 from dataclasses import dataclass
+from itertools import cycle
 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
@@ -18,7 +18,7 @@
  infer_num_classes,
  select_train_val_splits,
 )
-from .balanced_sampler import FewShotSampler
+from .balanced_sampler import BalancedSampler
 
 
 @dataclass
@@ -95,10 +95,12 @@ def load_prompts(
  Returns:
  An iterable dataset of prompts.
  """
+ class_counts = []
  prompters = []
- raw_datasets = []
+ datasets = []
  train_datasets = []
  rng = Random(seed)
+ assert num_shots == 0
 
  # First load the datasets and prompters. We need to know the minimum number of
  # templates for any dataset in order to make sure we don't run out of prompts.
@@ -112,29 +114,37 @@ def load_prompts(
  train_name, val_name = select_train_val_splits(ds_dict)
  split_name = val_name if split_type == "val" else train_name
 
- # Note that when streaming we can only approximately shuffle the dataset
- # using a buffer. Streaming shuffling is NOT an adequate shuffle for
- # datasets like IMDB, which are sorted by label.
- bad_streaming_datasets = ["imdb"]
- assert not (
- stream and ds_name in bad_streaming_datasets
- ), f"Streaming is not supported for {ds_name}."
- split = ds_dict[split_name].shuffle(seed=seed)
+ ds = ds_dict[split_name].shuffle(seed=seed)
  train_ds = ds_dict[train_name].shuffle(seed=seed)
+
  if not stream:
- split = assert_type(Dataset, split)
- split = split.to_iterable_dataset().cast(split.features)
+ ds = assert_type(Dataset, ds)
+ if world_size > 1:
+ ds = ds.shard(world_size, rank)
+
+ ds = ds.to_iterable_dataset().cast(ds.features)
 
- # only keep the datapoints relevant to the current process
- if world_size > 1:
+ elif world_size > 1:
  # This prints to stdout which is slightly annoying
- split = split_dataset_by_node(
- dataset=split, rank=rank, world_size=world_size
- )
+ ds = split_dataset_by_node(dataset=ds, rank=rank, world_size=world_size)
 
- raw_datasets.append(split)
+ label_column = infer_label_column(ds.features)
+ num_classes = infer_num_classes(ds.features[label_column])
+ if label_column != "label":
+ ds = ds.rename_column(label_column, "label")
+ train_ds = train_ds.rename_column(label_column, "label")
+
+ class_counts.append(num_classes)
+ datasets.append(ds)
  train_datasets.append(train_ds)
 
+ # Number of classes should be the same for all datasets
+ num_classes, *rest = class_counts
+ if not all(num_classes == x for x in rest):
+ raise ValueError(
+ f"# classes should be the same for all datasets, but got {class_counts}"
+ )
+
  min_num_templates = min(len(prompter.templates) for prompter in prompters)
  num_variants = (
  min_num_templates
@@ -145,51 +155,29 @@ def load_prompts(
  if rank == 0:
  print(f"Using {num_variants} variants of each prompt")
 
- ds_iterators = [iter(ds) for ds in raw_datasets]
- while True: # terminates when the first dataset runs out of examples
- for ds_iterator, ds, train_ds, prompter in zip(
- ds_iterators, raw_datasets, train_datasets, prompters
- ):
- label_column = infer_label_column(ds.features)
- num_classes = infer_num_classes(ds.features[label_column])
-
- # Remove everything except the label column
- extra_cols = list(assert_type(Features, ds.features))
- extra_cols.remove(label_column)
-
- if label_column != "label":
- ds = ds.rename_column(label_column, "label")
- if num_shots > 0:
- fewshot = FewShotSampler(
- train_ds, # TODO: not iterator
- num_shots=num_shots,
- rng=rng,
- )
- fewshot_iter = iter(fewshot)
- else:
- fewshot_iter = None
-
- try:
- example = next(ds_iterator)
- except StopIteration:
- return
-
- example = _convert_to_prompts(
- example,
- label_column=label_column,
- num_classes=num_classes,
- num_variants=num_variants,
- prompter=prompter,
- rng=rng,
- fewshot_iter=fewshot_iter,
- )
+ ds_iters = [iter(BalancedSampler(ds, num_classes)) for ds in datasets]
+ for ds_iter, ds, prompter in cycle(zip(ds_iters, datasets, prompters)):
+ try:
+ example = next(ds_iter)
+ except StopIteration:
+ return
+
+ example = _convert_to_prompts(
+ example,
+ label_column="label",
+ num_classes=num_classes,
+ num_variants=num_variants,
+ prompter=prompter,
+ rng=rng,
+ fewshot_iter=None,
+ )
 
-  # Add the builder and config name to the records directly to make
-  # sure we don't forget what dataset they came from.
-  example["builder_name"] = ds.info.builder_name
-  example["config_name"] = ds.info.config_name
+ # Add the builder and config name to the records directly to make
+ # sure we don't forget what dataset they came from.
+ example["builder_name"] = ds.info.builder_name
+ example["config_name"] = ds.info.config_name
 
-  yield example
+ yield example
 
 
 def _convert_to_prompts(

elk/metrics.py

@@ -1,7 +1,6 @@
 from functools import partial
 from typing import Literal
 
-import torch
 from sklearn.metrics import average_precision_score, roc_auc_score
 from torch import Tensor
 
@@ -17,7 +16,7 @@ def to_one_hot(labels: Tensor, n_classes: int) -> Tensor:
  Returns:
  Tensor: A one-hot representation tensor of shape (N, n_classes).
  """
- one_hot_labels = torch.zeros(labels.size(0), n_classes, dtype=torch.float32)
+ one_hot_labels = labels.new_zeros(labels.size(0), n_classes)
  return one_hot_labels.scatter_(1, labels.unsqueeze(1).long(), 1)
 
 

elk/training/ccs_reporter.py

@@ -83,12 +83,12 @@ class CcsReporter(Reporter):
 
  def __init__(
  self,
- in_features: int,
  cfg: CcsReporterConfig,
+ in_features: int,
  device: Optional[str] = None,
  dtype: Optional[torch.dtype] = None,
  ):
- super().__init__(in_features, cfg, device=device, dtype=dtype)
+ super().__init__(cfg, in_features, device=device, dtype=dtype)
 
  hidden_size = cfg.hidden_size or 4 * in_features // 3
 

elk/training/eigen_reporter.py

@@ -25,9 +25,9 @@ class EigenReporterConfig(ReporterConfig):
  of eigenvectors to compute from the VINC matrix.
  """
 
- var_weight: float = 1.0
- inv_weight: float = 5.0
- neg_cov_weight: float = 5.0
+ var_weight: float = 0.2
+ inv_weight: float = 1.0
+ neg_cov_weight: float = 1.0
 
  num_heads: int = 1
 
@@ -100,7 +100,6 @@ def forward(self, x: Tensor) -> Tensor:
 
  def predict(self, *hiddens: Tensor) -> Tensor:
  """Return the predicted logits on the contrast set `hiddens`."""
- # breakpoint()
  if len(hiddens) == 1:
  return self(hiddens[0])
 

elk/training/reporter.py

@@ -90,22 +90,22 @@ def __init__(
  @classmethod
  def check_separability(
  cls,
- train_pair: tuple[Tensor, Tensor],
- val_pair: tuple[Tensor, Tensor],
+ train_hiddens: Tensor,
+ val_hiddens: Tensor,
  ) -> float:
  """Measure how linearly separable the pseudo-labels are for a contrast pair.
 
  Args:
- train_pair: A tuple of tensors, (x0, x1), where x0 and x1 are the
+ train_hiddens: Tensor of shape [n, ], where x0 and x1 are the
  contrastive representations. Used for training the classifier.
  val_pair: A tuple of tensors, (x0, x1), where x0 and x1 are the
  contrastive representations. Used for evaluating the classifier.
 
  Returns:
  The AUROC of a linear classifier fit on the pseudo-labels.
  """
- x0, x1 = train_pair
- val_x0, val_x1 = val_pair
+ x0, x1 = train_hiddens
+ val_x0, val_x1 = val_hiddens
 
  pseudo_clf = Classifier(x0.shape[-1], device=x0.device) # type: ignore
  pseudo_train_labels = torch.cat(
@@ -198,13 +198,15 @@ def score(self, labels: Tensor, hiddens: Tensor) -> EvalResult:
  to_one_hot(Y, n_classes=c).long().flatten()
 
  if c == 2:
- cal_err = CalibrationError().update(Y.cpu(), pred_probs.cpu()).compute().ece
+ pos_probs = pred_probs[..., 0].flatten()
+ cal_err = CalibrationError().update(Y.cpu(), pos_probs.cpu()).compute().ece
+
  # Calibrated accuracy
- cal_thresh = pred_probs.float().quantile(labels.float().mean())
- cal_preds = pred_probs.gt(cal_thresh).squeeze(1).to(torch.int)
+ cal_thresh = pos_probs.float().quantile(labels.float().mean())
+ cal_preds = pos_probs.gt(cal_thresh).to(torch.int)
  cal_acc = cal_preds.flatten().eq(Y).float().mean().item()
 
- raw_preds = pred_probs.gt(0.5).squeeze(1).to(torch.int)
+ raw_preds = pos_probs.gt(0.5).to(torch.int)
  else:
  # TODO: Implement calibration error for k > 2?
  cal_acc = 0.0

elk/training/train.py

@@ -90,14 +90,14 @@ def train_reporter(
  # pseudo_auroc = self.get_pseudo_auroc(layer, x0, x1, val_x0, val_x1)
 
  if isinstance(self.cfg.net, CcsReporterConfig):
- reporter = CcsReporter(d, self.cfg.net, device=device)
+ reporter = CcsReporter(self.cfg.net, d, device=device)
  elif isinstance(self.cfg.net, EigenReporterConfig):
  reporter = EigenReporter(self.cfg.net, d, c, device=device)
  else:
  raise ValueError(f"Unknown reporter config type: {type(self.cfg.net)}")
 
  train_loss = reporter.fit(*train_h.unbind(2), labels=train_gt)
- val_result = reporter.score(val_gt, val_h)
+ val_result = reporter.score(val_gt.to(device), val_h)
 
  reporter_dir, lr_dir = self.create_models_dir(assert_type(Path, self.out_dir))
  if val_lm_preds is not None: