|
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