Make load_prompts actually use label_column and num_classes args

michaelbyun · Apr 18, 2023 · e1b1834 · e1b1834
1 parent 54f5ed3
commit e1b1834
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/elk/extraction/prompt_loading.py b/elk/extraction/prompt_loading.py
@@ -160,8 +160,8 @@ def load_prompts(
  if rank == 0:
  print(f"Using {num_variants} variants of each prompt")
 
- label_column = infer_label_column(ds.features)
- num_classes = infer_num_classes(ds.features[label_column])
+ label_column = label_column or infer_label_column(ds.features)
+ num_classes = num_classes or infer_num_classes(ds.features[label_column])
  rng = Random(seed)
 
  if num_shots > 0: