Add files via upload

easyjailbreak/selector/UCBSelectPolicy.py

@@ -0,0 +1,60 @@
+"""
+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)