github/privateGPT
1
0
Fork 1
mirror of https://github.com/imartinez/privateGPT.git synced 2025-09-23 20:17:24 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: fix tests

Browse Source
This commit is contained in:
Pablo Orgaz
2024-05-30 15:37:23 +02:00
parent 4e99457ea5
commit e37d3b9235
2 changed files with 19 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
private_gpt/components/llm/prompt_helper.py
View File

@@ -176,16 +176,18 @@ class MistralPromptStyle(AbstractPromptStyle):
inst_buffer = []
text = ""
for message in messages:
if message.role == MessageRole.SYSTEM:
if message.role == MessageRole.SYSTEM or message.role == MessageRole.USER:
inst_buffer.append(str(message.content).strip())
elif message.role == MessageRole.USER:
inst_buffer.append(str(message.content).strip())
text += "<s>[/INST] " + "\n".join(inst_buffer) + " [/INST]"
inst_buffer.clear()
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:

20
tests/test_prompt_helper.py
View File

@@ -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()