Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

fix: mistral ignoring assistant messages #1954

Merged
merged 3 commits into from
May 30, 2024
Merged

fix: mistral ignoring assistant messages #1954

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
136f8b5
fix: mistral ignoring assistant messages
pabloogc May 28, 2024
4e99457
fix: typing
pabloogc May 29, 2024
e37d3b9
fix: fix tests
pabloogc May 30, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 15 additions & 11 deletions private_gpt/components/llm/prompt_helper.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -173,18 +173,22 @@ def _completion_to_prompt(self, completion: str) -> str:

class MistralPromptStyle(AbstractPromptStyle):
def _messages_to_prompt(self, messages: Sequence[ChatMessage]) -> str:
prompt = "<s>"
inst_buffer = []
text = ""
for message in messages:
role = message.role
content = message.content or ""
if role.lower() == "system":
message_from_user = f"[INST] {content.strip()} [/INST]"
prompt += message_from_user
elif role.lower() == "user":
prompt += "</s>"
message_from_user = f"[INST] {content.strip()} [/INST]"
prompt += message_from_user
return prompt
if message.role == MessageRole.SYSTEM or message.role == MessageRole.USER:
inst_buffer.append(str(message.content).strip())
elif message.role == MessageRole.ASSISTANT:
text += "<s>[INST] " + "\n".join(inst_buffer) + " [/INST]"
text += " " + str(message.content).strip() + "</s>"
inst_buffer.clear()
else:
raise ValueError(f"Unknown message role {message.role}")

if len(inst_buffer) > 0:
text += "<s>[INST] " + "\n".join(inst_buffer) + " [/INST]"

return text

def _completion_to_prompt(self, completion: str) -> str:
return self._messages_to_prompt(
Expand Down
20 changes: 12 additions & 8 deletions tests/test_prompt_helper.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,17 +69,21 @@ def test_tag_prompt_style_format_with_system_prompt():
def test_mistral_prompt_style_format():
prompt_style = MistralPromptStyle()
messages = [
ChatMessage(content="You are an AI assistant.", role=MessageRole.SYSTEM),
ChatMessage(content="Hello, how are you doing?", role=MessageRole.USER),
ChatMessage(content="A", role=MessageRole.SYSTEM),
ChatMessage(content="B", role=MessageRole.USER),
]

expected_prompt = (
"<s>[INST] You are an AI assistant. [/INST]</s>"
"[INST] Hello, how are you doing? [/INST]"
)

expected_prompt = "<s>[INST] A\nB [/INST]"
assert prompt_style.messages_to_prompt(messages) == expected_prompt

messages2 = [
ChatMessage(content="A", role=MessageRole.SYSTEM),
ChatMessage(content="B", role=MessageRole.USER),
ChatMessage(content="C", role=MessageRole.ASSISTANT),
ChatMessage(content="D", role=MessageRole.USER),
]
expected_prompt2 = "<s>[INST] A\nB [/INST] C</s><s>[INST] D [/INST]"
assert prompt_style.messages_to_prompt(messages2) == expected_prompt2


def test_chatml_prompt_style_format():
prompt_style = ChatMLPromptStyle()
Expand Down
Loading