Merge pull request #122 from EleutherAI/sort-variant

Sort outputs by variant id during extraction
EleutherAI · Mar 17, 2023 · 3af509b · 3af509b
2 parents 1bb7f1e + 0f78b9a
commit 3af509b
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/elk/extraction/prompt_dataset.py b/elk/extraction/prompt_dataset.py
@@ -179,8 +179,10 @@ def __init__(
  def __getitem__(self, index: int) -> list[Prompt]:
  """Get a list of prompts for a given predicate"""
  # get self.num_variants unique prompts from the template pool
- template_names = self.rng.sample(
- list(self.prompter.templates), self.num_variants
+ template_names = (
+ self.rng.sample(list(self.prompter.templates), self.num_variants)
+ if self.num_variants < len(self.prompter.templates)
+ else list(self.prompter.templates)
  )
 
  example = self.active_split[index]