Add files via upload

easyjailbreak/mutation/generation/ApplyGPTMutation.py

@@ -0,0 +1,73 @@
+import openai
+import random
+import sys
+import time
+from easyjailbreak.mutation import MutationBase
+
+class ApplyGPTMutation(MutationBase):
+    r"""
+    The 'ApplyGPTMutation' class, inheriting from 'MutationBase', is designed to apply
+    mutations to text instances using GPT-generated modifications. This class specifically
+    utilizes OpenAI's GPT models to mutate text based on a given mutation rate.
+    """
+
+    def __init__(self, attr_name="attack_attrs", API_key=None, mutation_rate=0.01):
+        r"""
+        Initializes the ApplyGPTMutation instance with an API key, attribute name, and mutation rate.
+        :param str attr_name: The attribute name in the instance where the candidate prompts are stored.
+        :param str|None API_key: The API key for accessing OpenAI's GPT models.
+        :param float mutation_rate: The rate at which mutations are applied to text instances.
+        """
+        self.API_key = API_key
+        self.attr_name = attr_name
+        self.mutation_rate = mutation_rate
+        
+    def _get_mutated_instance(self, instance, **kwargs):
+        r"""
+         Generates a mutated instance of the given object using GPT model-based mutations.
+         :param instance: The instance to be mutated.
+         :param kwargs: Additional keyword arguments.
+         :return: A list containing the mutated instance.
+         """
+        mutated_prompts = []
+        seeds = getattr(instance, self.attr_name)["candidate_prompts"]
+        for seed in seeds:
+            if random.random() < self.mutation_rate:
+                mutated_prompts.append(self.gpt_mutate(seed, self.API_key))
+        new_instance = instance.copy()
+        # setattr(new_instance, self.attr_name, mutated_prompts)
+        new_instance.attack_attrs["candidate_prompts"] = mutated_prompts
+        return [new_instance]
+
+    def gpt_mutate(self, sentence, API_key):
+        r"""
+        Applies GPT model-based mutations to a given sentence.
+        :param str sentence: The sentence to be mutated.
+        :param str API_key: The API key for accessing the GPT model.
+        :return str: The mutated version of the sentence.
+        """
+        client = openai.OpenAI(api_key=API_key)
+        system_msg = 'You are a helpful and creative assistant who writes well.'
+        user_message = f"Please revise the following sentence with no changes to its length and only output the revised version, the sentences are: \n '{sentence}'."
+        revised_sentence = sentence
+        received = False
+        while not received:
+            try:
+                response = client.chat.completions.create(model="gpt-4",
+                                                        messages=[{"role": "system", "content": system_msg},
+                                                              {"role": "user", "content": user_message}],
+                                                        temperature=1, top_p=0.9)
+                revised_sentence = response.choices[0].message.content.replace('\n', '')
+                received = True
+            except:
+                error = sys.exc_info()[0]
+                print("error:", error)
+                time.sleep(1)
+        if revised_sentence.startswith("'") or revised_sentence.startswith('"'):
+            revised_sentence = revised_sentence[1:]
+        if revised_sentence.endswith("'") or revised_sentence.endswith('"'):
+            revised_sentence = revised_sentence[:-1]
+        if revised_sentence.endswith("'.") or revised_sentence.endswith('".'):
+            revised_sentence = revised_sentence[:-2]
+        print(f'revised: {revised_sentence}')
+        return revised_sentence