fix label assignment causing multiclass bug (EleutherAI#183)

michaelbyun · Apr 12, 2023 · fe3a014 · fe3a014
1 parent 90282c4
commit fe3a014
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/elk/extraction/prompt_loading.py b/elk/extraction/prompt_loading.py
@@ -203,7 +203,7 @@ def _convert_to_prompts(
  fewshot_iter: Optional[Iterator[list[dict]]] = None,
 ) -> dict[str, Any]:
  """Prompt-generating function to pass to `IterableDataset.map`."""
- label = assert_type(int, example[label_column])
+ assert_type(int, example[label_column])
  prompts = []
  templates = list(prompter.templates.values())
  if num_variants < len(templates):
@@ -259,7 +259,7 @@ def qa_cat(q: str, a: str) -> str:
  raise ValueError(f'Prompt duplicated {dup_count} times! "{maybe_dup}"')
 
  return dict(
- label=label,
+ label=new_label,
  prompts=prompts,
  template_names=prompter.all_template_names,
  )