Add --no_balance flag to not balance datasets

elk/debug_logging.py

@@ -31,6 +31,9 @@ def save_debug_log(datasets: list[DatasetDictWithName], out_dir: Path) -> None:
         else:
             train_split, val_split = select_train_val_splits(ds)
 
+        if len(ds[val_split]) == 0:
+            logging.warning(f"Val split '{val_split}' is empty!")
+            continue
         text_questions = ds[val_split][0]["text_questions"]
         template_ids = ds[val_split][0]["variant_ids"]
         label = ds[val_split][0]["label"]

elk/extraction/extraction.py

@@ -65,6 +65,9 @@ class Extract(Serializable):
     binarize: bool = False
     """Whether to binarize the dataset labels for multi-class datasets."""
 
+    no_balance: bool = False
+    """Whether to disable balancing the dataset by label."""
+
     int8: bool = False
     """Whether to perform inference in mixed int8 precision with `bitsandbytes`."""
 
@@ -189,6 +192,7 @@ def extract_hiddens(
         num_shots=cfg.num_shots,
         split_type=split_type,
         template_path=cfg.template_path,
+        balance=not cfg.no_balance,
         rank=rank,
         world_size=world_size,
         seed=cfg.seed,

elk/extraction/prompt_loading.py

@@ -21,6 +21,7 @@ def load_prompts(
     seed: int = 42,
     split_type: Literal["train", "val"] = "train",
     template_path: str | None = None,
+    balance: bool = True,
     rank: int = 0,
     world_size: int = 1,
 ) -> Iterator[dict]:
@@ -89,14 +90,14 @@ def load_prompts(
     else:
         fewshot_iter = None
 
-    if label_column in ds.features:
+    if label_column in ds.features and balance:
         ds = BalancedSampler(
             ds.to_iterable_dataset(),
             set(label_choices),
             label_col=label_column,
         )
     else:
-        if rank == 0:
+        if rank == 0 and balance:
             print("No label column found, not balancing")
         ds = ds.to_iterable_dataset()
 

elk/training/classifier.py

@@ -194,8 +194,8 @@ def inlp(
                 the input dimension.
             y: Target tensor of shape (N,) for binary classification or (N, C) for
                 multiclass classification, where C is the number of classes.
-            max_iter: Maximum number of iterations to run. If `None`, run for the full
-                dimension of the input.
+            max_iter: Maximum number of iterations to run. If `None`, run until the data
+                is linearly guarded (no linear classifier can extract information).
             tol: Tolerance for the loss function. The algorithm will stop when the loss
                 is within `tol` of the entropy of the labels.
 
@@ -212,12 +212,11 @@ def inlp(
         p = y.float().mean()
         H = -p * torch.log(p) - (1 - p) * torch.log(1 - p)
 
-        if max_iter is not None:
-            d = min(d, max_iter)
+        max_iter = max_iter or d
 
         # Iterate until the loss is within epsilon of the entropy
         result = InlpResult()
-        for _ in range(d):
+        for _ in range(max_iter):
             clf = cls(d, device=x.device, dtype=x.dtype)
             loss = clf.fit(x, y)
             result.classifiers.append(clf)

elk/training/supervised.py

@@ -6,7 +6,7 @@
 
 
 def train_supervised(
-    data: dict[str, tuple], device: str, mode: str
+    data: dict[str, tuple], device: str, mode: str, max_inlp_iter: int | None = None
 ) -> list[Classifier]:
     Xs, train_labels = [], []
 
@@ -26,7 +26,7 @@ def train_supervised(
         lr_model.fit_cv(X, train_labels)
         return [lr_model]
     elif mode == "inlp":
-        return Classifier.inlp(X, train_labels).classifiers
+        return Classifier.inlp(X, train_labels, max_inlp_iter).classifiers
     elif mode == "single":
         lr_model = Classifier(X.shape[-1], device=device)
         lr_model.fit(X, train_labels)

elk/training/train.py

@@ -34,6 +34,9 @@ class Elicit(Run):
     cross-validation. Defaults to "single", which means to train a single classifier
     on the training data. "cv" means to use cross-validation."""
 
+    max_inlp_iter: int | None = None
+    """Maximum number of iterations for Iterative Nullspace Projection (INLP)."""
+
     def create_models_dir(self, out_dir: Path):
         lr_dir = None
         lr_dir = out_dir / "lr_models"
@@ -124,6 +127,7 @@ def apply_to_layer(
                 train_dict,
                 device=device,
                 mode=self.supervised,
+                max_inlp_iter=self.max_inlp_iter,
             )
             with open(lr_dir / f"layer_{layer}.pt", "wb") as file:
                 torch.save(lr_models, file)