# this file contains various prompts for various models/datasets

PROMPT_COLLIE_PHI_INSTRUCT_CUSTOMIZED = \
"""<|system|>
You are an assistant that creates text according to provided requirements.\
 Read the users requirements carefully and compose the output text as specified. \
 Only provide the text that fits the given criteria and nothing else.\
 DO NOT PROVIDE quotation marks.\
 DO NOT PROVIDE assertions (e.g. 'Here is a sentence with' or 'Here is a paragraph that' and similar).\
 DO NOT PROVIDE confirmations or summaries (e.g. 'the number of characters in this text is ...')!
 The correctness of the answers will be checked automatically and such unnecessary things will fail the otherwise correct response!\
<|end|>
<|user|>
{question}<|end|>
<|assistant|>"""



PROMPT_COLLIE_PHI_INSTRUCT_CUSTOMIZED_ONESHOT = \
"""<|system|>
You are an assistant that creates text according to provided requirements.\
 Read the users requirements carefully and compose the output text as specified. \
 Only provide the text that fits the given criteria and nothing else.\
 DO NOT GENERATE additional quotation marks or additional explanations or assertions (e.g. 'Here is a sentence with' or 'Here is a paragraph that' and similar) or confirmations (e.g. 'yep, the number of characters is right!')!
 The correctness of the answers will be checked automatically and such unnecessary things will faile the otherwise correct response!\
<|end|>
<|user|>
{example_in}<|end|>
<|assistant|>
{example_out}<|end|>
<|user|>
{question}<|end|>
<|assistant|>"""


PROMPT_COLLIE_LLAMA_INSTRUCT_CUSTOMIZED = \
"""
<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n\
You are an assistant that creates text according to provided requirements. \
 Read the users requirements carefully and compose the output text as specified. \
 Only provide the text that fits the given criteria and nothing else.\
 DO NOT PROVIDE quotation marks.\
 DO NOT PROVIDE assertions (e.g. 'Here is a sentence with' or 'Here is a paragraph that' and similar).\
 DO NOT PROVIDE confirmations or summaries (e.g. 'the number of characters in this text is ...')!
 The correctness of the answers will be checked automatically and such unnecessary things will fail the otherwise correct response!\
<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n\
{question}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n\
"""

PROMPT_COLLIE_LLAMA_INSTRUCT_CUSTOMIZED_ONESHOT = \
"""
<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n\
You are an assistant that creates text according to provided requirements. 
Read the users requirements carefully and compose the output text as specified. \
Only provide the text that fits the given criteria and nothing else (including no additional quotation marks or additional explanations or confirmations like 'yep, the number of characters is right!') as its correctness will be checked automatically and such things will throw off the checkers.\
Ignore the topic of the text in the example given and focus on fulfilling constraints.\
<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n\
{example_in}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n\
{example_out}<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n\
{question}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n\
"""


PROMPT_KUQ_PHI_INSTRUCT_CUSTOMIZED = \
"""<|system|>
Read the following question carefully and answer it.\
 Think before answering.\
 Give specific and concise answers to the questions, preferably no more than a couple of sentences.\
 Only answering the given question is required.\
<|end|>
<|user|>
### Question: {question}<|end|>
<|assistant|>
### Answer: """


PROMPT_KUQ_LLAMA_INSTRUCT_CUSTOMIZED = \
"""<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n\
Read the following question carefully and answer it.\
 Think before answering.\
 Give specific and concise answers to the questions, preferably no more than a couple of sentences.\
 Only answering the given question is required.\
<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n\
### Question: {question}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n\
### Answer: \
"""


PROMPT_COQA_PHI_35_INSTRUCT = \
"""<|system|>
You are an intelligent question answering assistant that gives answers to questions relating a given text.\
 Read the following text carefully and answer the users question concisely:\n\n{background}<|end|>
<|user|>
### Question: {question}<|end|>
<|assistant|>
### Answer: """

PROMPT_COQA_LLAMA_3_INSTRUCT = \
"""<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n\
You are an intelligent question answering assistant that gives answers to questions relating a given text.\
 Read the following text carefully and answer the users question concisely:\n\n\
{background}<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n\
### Question: {question}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n\
### Answer: \
"""


PROMPT_TRIVIA_PHI_35_INSTRUCT = \
"""<|system|>
You are an intelligent question answering assistant that gives concise answers to questions asked by the user.<|end|>
<|user|>
### Question: {question}<|end|>
<|assistant|>
### Answer: """

PROMPT_TRIVIA_LLAMA_3_INSTRUCT = \
"""<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n\
You are an intelligent question answering assistant that gives answers to questions asked by the user.\
<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n\
### Question: {question}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n\
### Answer: \
"""

PROMPT_NO_INSTRUCT_MODEL = \
"""{question}"""



#### DO NOT MODIFY!!!
PROMPT_PHI_35_INSTRUCT_GENERAL = \
"""<|system|>
{system_message}<|end|>
<|user|>
{question}<|end|>
<|assistant|>"""

PROMPT_LLAMA_3_GENERAL_TEMPLATE = \
"""<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n\
{{ system_prompt }}<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n\
{{ user_msg_1 }}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n\
{{ model_answer_1 }}<|eot_id|>\
"""


from abc import abstractmethod
from collections import defaultdict


class PromptBuilder():
    def __init__(self, role_prefix_mapping={}):
        self.rpm = defaultdict(str, role_prefix_mapping)

    @abstractmethod
    def wrap_message(self, role, message, last_msg=False):
        pass

    @abstractmethod
    def finalize(self, final_prompt):
        pass

    def get_system_role(self):
        return 'system'

    def get_user_role(self):
        return 'user'

    def get_assistant_role(self):
        return 'assistant'

    def build_prompt(self,chat_sequence,starting_string=''):
        final_prompt = starting_string
        for msg in chat_sequence:
            final_prompt += self.wrap_message(**msg)
        return self.finalize(final_prompt)


class PromptBuilderLlama3I(PromptBuilder):
    def wrap_message(self, role, message, last_msg=False):
        return f'<|start_header_id|>{role}<|end_header_id|>\n\n{self.rpm[role]}{message}' + ('<|eot_id|>' if not last_msg else ' ')

    def finalize(self, final_prompt):
        return '<|begin_of_text|>' + final_prompt



class PromptBuilderQwen25(PromptBuilder):
    def wrap_message(self, role, message, last_msg=False):
        if role=='system':
            message = "You are Qwen, created by Alibaba Cloud. You are a helpful assistant." + message
            
        return f'"<|im_start|>{role}\n{self.rpm[role]}{message}' + ('<|im_end|>\n' if not last_msg else ' ') + '\n'

    def finalize(self, final_prompt):
        return final_prompt # no bos tokens for qwen



class PromptBuilderPhi35I(PromptBuilder):
    def wrap_message(self, role, message, last_msg=False):
        assert role in ['system', 'user', 'assistant'], f"Misspecified role {role}!"
        return f'<|{role}|>\n{self.rpm[role]}{message}' + ('<|end|>\n' if not last_msg else ' ')

    def finalize(self, final_prompt):
        return final_prompt


class PromptBuilderNonIFT(PromptBuilder):
    def wrap_message(self, role, message, last_msg=False):
        return ((role+"\n") if role!="system" else "") + self.rpm[role] + message + ('\n\n' if not last_msg else ' ')

    def finalize(self, final_prompt):
        return final_prompt


class PromptBuilderLlama3ForBCB(PromptBuilder):
    # non instruction tuned llama, no extras for BCB code instruct template
    # IMPORTANT: the code actually using this must respect user roles!
    def get_system_role(self):
        return ''
    
    def get_user_role(self):
        return ''

    def get_assistant_role(self):
        return ''

    def wrap_message(self, role, message, last_msg=True):
        return role+self.rpm[role]+message

    def finalize(self, final_prompt):
        return final_prompt


class PromptBuilderFalcon3(PromptBuilder):
    # "{{bos_token}}{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}"
    def wrap_message(self, role, message, last_msg=False):
        assert role in ['system', 'user', 'assistant'], f"Misspecified role {role}!"
        return f"<|im_start|>{role}\n{self.rpm[role]}{message}" + (f"<im_end>\n" if not last_msg else ' ')
        # f'<|{role}|>\n{self.rpm[role]}{message}' + ('<|end|>\n' if not last_msg else '')

    def finalize(self, final_prompt):
        return '<|begin_of_text|>'+final_prompt


def get_prompt_builder_by_model_id(model_id: str, dataset=None):
    if dataset=='BCB':
        # return the blanked out prompt builder, otherwise will break the code generation
        return PromptBuilderLlama3ForBCB()
    elif model_id.startswith("meta-llama/Meta-Llama-3-"):
        return PromptBuilderLlama3I()
    elif model_id == "microsoft/Phi-3" or model_id.startswith("microsoft/Phi-3"):
        return PromptBuilderPhi35I()
    elif model_id == "Zyphra/Zamba-7B-v1":
        return PromptBuilderNonIFT()
    elif model_id.startswith("tiiuae/Falcon3"):
        return PromptBuilderFalcon3()
    elif 'r1' in model_id:
        return PromptBuilderR1()
    elif 'Qwen2.5' in model_id:
        return PromptBuilderQwen25()
    else:
        raise AttributeError(f"No prompt builder for {model_id}!")



## R1 prompter
# User: prompt. Assistant:
class PromptBuilderR1(PromptBuilder):
    def __init__(self, role_prefix_mapping={}, think_assistant=True):
        super(PromptBuilderR1, self).__init__(role_prefix_mapping=role_prefix_mapping)
        self.think_assistant = think_assistant

    def get_system_role(self):
        return 'unused!'
    
    def get_user_role(self):
        return 'User'

    def get_assistant_role(self):
        return 'Assistant'

    def wrap_message(self, role, message, last_msg=False):
        if role == self.get_system_role():
            # no system prompts for r1
            return f"{message}"
        else:
            main = f' {role}: {self.rpm[role]}{message}'
            suffix = ''
            if not last_msg and not (message.endswith('.') or message.endswith('\n')):
                suffix += '. '
            elif last_msg and self.think_assistant:
                suffix += ' <think>'
            return main + suffix

    def finalize(self, final_prompt):
        return final_prompt


class PromptBuilderLlaSMol(PromptBuilder):
    # they don't seem to use any system prompts in their finetunes ???
    def get_system_role(self):
        return 'unused!'
    
    def get_user_role(self):
        return 'Query'

    def get_assistant_role(self):
        return 'Response'
    
    def wrap_message(self, role, message, last_msg=False):
        if role == self.get_system_role():
            return ""
        return f'{role} :' + (f' {self.rpm[role]}{message}\n\n' if not last_msg else '')

    def finalize(self, final_prompt):
        return '<s>' + final_prompt


# class PromptBuilderQWEN25(PromptBuilder):
#     # they don't seem to use any system prompts in their finetunes ???
#     def get_system_role(self):
#         return 'system'
    
#     def get_user_role(self):
#         return 'user'

#     def get_assistant_role(self):
#         return 'assistant'
    
#     def wrap_message(self, role, message, last_msg=False):
#         return f'<|im_start|>{role}\n' + self.rpm[role] + message + ('<|im_end|>\n' if not last_msg else '')

#     def finalize(self, final_prompt):
#         return final_prompt


