only preserve order for the all variants case

elk/extraction/extraction.py

@@ -163,13 +163,6 @@ def collate(prompts: list[Prompt]) -> list[list[BatchEncoding]]:
     # Iterating over questions
     layer_indices = cfg.layers or tuple(range(model.config.num_hidden_layers))
     for prompts in prompt_ds:
-        variant_ids = [prompt.template_name for prompt in prompts]
-
-        # Sort the variants and prompts by their ID to standardize the order
-        sorted_idxs = sorted(range(len(variant_ids)), key=variant_ids.__getitem__)
-        variant_ids = [variant_ids[i] for i in sorted_idxs]
-        prompts = [prompts[i] for i in sorted_idxs]
-
         inputs = collate(prompts)
         hidden_dict = {
             f"hidden_{layer_idx}": torch.empty(
@@ -181,6 +174,7 @@ def collate(prompts: list[Prompt]) -> list[list[BatchEncoding]]:
             )
             for layer_idx in layer_indices
         }
+        variant_ids = [prompt.template_name for prompt in prompts]
 
         # Iterate over variants
         for i, variant_inputs in enumerate(inputs):

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]