flag

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()