from conversation import Conversation
from models.base import BaseModel


class ChatML(BaseModel):

    def _get_prompt(self, conversation: Conversation):
        system_message = "\n".join(
            [conversation.memory, conversation.prompt]
        ).strip()
        prompt = f"<|im_start|>system\n{system_message}<|im_end|>"
        for message in conversation.messages:
            prompt += f"\n<|im_start|>{message['from']}\n{message['value']}<|im_end|>"
        prompt += f"\n<|im_start|>{conversation.bot_label}\n"
        return prompt