| 43 | end_token: str = "<|end|>" |
| 44 | |
| 45 | def get_training_prompt(self) -> str: |
| 46 | prompt = self.system_token + "\n" + self.system + self.end_token + "\n" |
| 47 | if self.messages is None: |
| 48 | raise ValueError("Dialogue template must have at least one message.") |
| 49 | for message in self.messages: |
| 50 | if message["role"] == "user": |
| 51 | prompt += self.user_token + "\n" + message["content"] + self.end_token + "\n" |
| 52 | else: |
| 53 | prompt += self.assistant_token + "\n" + message["content"] + self.end_token + "\n" |
| 54 | return prompt |
| 55 | |
| 56 | def get_inference_prompt(self) -> str: |
| 57 | prompt = self.system_token + "\n" + self.system + self.end_token + "\n" |