60 lines
2.5 KiB
Python
60 lines
2.5 KiB
Python
"""
|
|
UCBSelectPolicy class
|
|
==========================
|
|
"""
|
|
import numpy as np
|
|
from easyjailbreak.datasets import JailbreakDataset
|
|
from easyjailbreak.selector.selector import SelectPolicy
|
|
|
|
__all__ = ["UCBSelectPolicy"]
|
|
|
|
class UCBSelectPolicy(SelectPolicy):
|
|
"""
|
|
A selection policy based on the Upper Confidence Bound (UCB) algorithm. This policy is designed
|
|
to balance exploration and exploitation when selecting instances from a JailbreakDataset.
|
|
It uses the UCB formula to select instances that either have high rewards or have not been explored much.
|
|
"""
|
|
def __init__(self,
|
|
explore_coeff: float = 1.0,
|
|
Dataset: JailbreakDataset = None):
|
|
"""
|
|
Initializes the UCBSelectPolicy with a given JailbreakDataset and exploration coefficient.
|
|
|
|
:param float explore_coeff: Coefficient to control the exploration-exploitation balance.
|
|
:param ~JailbreakDataset Dataset: The dataset from which instances will be selected.
|
|
"""
|
|
super().__init__(Dataset)
|
|
|
|
self.step = 0
|
|
self.last_choice_index:int = 0
|
|
self.explore_coeff = explore_coeff
|
|
self.rewards = [0 for _ in range(len(Dataset))]
|
|
|
|
def select(self) -> JailbreakDataset:
|
|
"""
|
|
Selects an instance from the dataset based on the UCB algorithm.
|
|
|
|
:return ~JailbreakDataset: The selected JailbreakDataset from the dataset.
|
|
"""
|
|
if len(self.Datasets) > len(self.rewards):
|
|
self.rewards.extend([0 for _ in range(len(self.Datasets) - len(self.rewards))])
|
|
|
|
self.step += 1
|
|
scores = np.zeros(len(self.Datasets))
|
|
for i, prompt_node in enumerate(self.Datasets):
|
|
smooth_visited_num = prompt_node.visited_num + 1
|
|
scores[i] = self.rewards[i] / smooth_visited_num + \
|
|
self.explore_coeff * np.sqrt(2 * np.log(self.step) / smooth_visited_num)
|
|
|
|
self.last_choice_index = int(np.argmax(scores))
|
|
self.Datasets[self.last_choice_index].visited_num += 1
|
|
return JailbreakDataset([self.Datasets[self.last_choice_index]])
|
|
|
|
def update(self, Dataset: JailbreakDataset):
|
|
"""
|
|
Updates the rewards for the last selected instance based on the success of the prompts.
|
|
|
|
:param ~JailbreakDataset Dataset: The dataset containing prompts used for updating rewards.
|
|
"""
|
|
succ_num = sum([prompt_node.num_jailbreak for prompt_node in Dataset])
|
|
self.rewards[self.last_choice_index] += succ_num / len(Dataset)
|