File size: 2,379 Bytes
aa73958 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 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 |