|
7 | 7 | import logging
|
8 | 8 | import random
|
9 | 9 | from typing import Any, Callable, Dict, List, Literal, Optional, Union, cast
|
| 10 | +from itertools import zip_longest |
10 | 11 |
|
11 | 12 | from tqdm import tqdm
|
12 | 13 |
|
@@ -215,17 +216,18 @@ async def __call__(
|
215 | 216 | ncols=100,
|
216 | 217 | unit="simulations",
|
217 | 218 | )
|
218 |
| - for template in templates: |
219 |
| - parameter_order = list(range(len(template.template_parameters))) |
220 |
| - if randomize_order: |
221 |
| - # The template parameter lists are persistent across sim runs within a session, |
222 |
| - # So randomize a the selection instead of the parameter list directly, |
223 |
| - # or a potentially large deep copy. |
224 |
| - if randomization_seed is not None: |
225 |
| - random.seed(randomization_seed) |
226 |
| - random.shuffle(parameter_order) |
227 |
| - for index in parameter_order: |
228 |
| - parameter = template.template_parameters[index].copy() |
| 219 | + |
| 220 | + if randomize_order: |
| 221 | + # The template parameter lists are persistent across sim runs within a session, |
| 222 | + # So randomize a the selection instead of the parameter list directly, |
| 223 | + # or a potentially large deep copy. |
| 224 | + if randomization_seed is not None: |
| 225 | + random.seed(randomization_seed) |
| 226 | + random.shuffle(templates) |
| 227 | + parameter_lists = [t.template_parameters for t in templates] |
| 228 | + zipped_parameters = list(zip_longest(*parameter_lists)) |
| 229 | + for param_group in zipped_parameters: |
| 230 | + for template, parameter in zip(templates, param_group): |
229 | 231 | if _jailbreak_type == "upia":
|
230 | 232 | parameter = self._join_conversation_starter(parameter, random.choice(jailbreak_dataset))
|
231 | 233 | tasks.append(
|
|
0 commit comments