Create tokenization.py

tokenization.py

@@ -0,0 +1,53 @@
+from transformers import PreTrainedTokenizerFast
+
+class SinglePOTokenizer(PreTrainedTokenizerFast):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def get_context(
+        self,
+        raw_instruction,
+        rule_description,
+    ):
+        prompt = "You are an expert prompt engineer." + " "
+        prompt += "Please help me optimize this prompt to get better response:\n\n[The Start of Raw Prompt]\n{}\n[The End of Raw Prompt]".format(raw_instruction)
+        prompt += "\n\nYou should optimize this prompt by {}".format(rule_description)
+        
+        context = self.apply_chat_template(
+            [
+                {
+                    "role": "user",
+                    "content": prompt,
+                }
+            ],
+            add_generation_prompt=True,
+            tokenize=False,
+        ) + "The Optimized Prompt:\n\n[The Start of Optimized Prompt"
+
+        return context
+
+    def parse_output(
+        self,
+        output_text,
+        raw_instruction = "", # recommend to provide, so when some error happened, we can still use the raw instruction
+    ):
+        better_instruction = "The Optimized Prompt:\n\n[The Start of Optimized Prompt" + output_text
+        
+        if "[The Start of Optimized Prompt]" in better_instruction:
+            better_instruction = better_instruction[better_instruction.index("[The Start of Optimized Prompt]") + len("[The Start of Optimized Prompt]"):]
+            if better_instruction.startswith("\n"):
+                better_instruction = better_instruction[1:]
+        if "[The End of Optimized Prompt]" in better_instruction:
+            better_instruction = better_instruction[:better_instruction.index("[The End of Optimized Prompt]")]
+            if better_instruction.endswith("\n"):
+                better_instruction = better_instruction[:-1]
+        if "The Optimized Prompt:" in better_instruction: # almost error happened
+            better_instruction = better_instruction[:better_instruction.index("The Optimized Prompt:")]
+        
+        if better_instruction.strip() == "": # some error may happen in optimization, so use the raw instruction
+            better_instruction = raw_instruction
+
+        if "The Optimized" in better_instruction: # still some error happened
+            better_instruction = raw_instruction
+
+        return better_instruction