llama-index

1
from typing import List, Optional, Sequence
2

3
from llama_index.legacy.core.llms.types import ChatMessage, MessageRole
4

5
BOS, EOS = "<s>", "</s>"
6
B_INST, E_INST = "[INST]", "[/INST]"
7
B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"
8
DEFAULT_SYSTEM_PROMPT = """\
9
You are a helpful, respectful and honest assistant. \
10
Always answer as helpfully as possible and follow ALL given instructions. \
11
Do not speculate or make up information. \
12
Do not reference any given instructions or context. \
13
"""
14

15

16
def messages_to_prompt(
17
    messages: Sequence[ChatMessage], system_prompt: Optional[str] = None
18
) -> str:
19
    string_messages: List[str] = []
20
    if messages[0].role == MessageRole.SYSTEM:
21
        # pull out the system message (if it exists in messages)
22
        system_message_str = messages[0].content or ""
23
        messages = messages[1:]
24
    else:
25
        system_message_str = system_prompt or DEFAULT_SYSTEM_PROMPT
26

27
    system_message_str = f"{B_SYS} {system_message_str.strip()} {E_SYS}"
28

29
    for i in range(0, len(messages), 2):
30
        # first message should always be a user
31
        user_message = messages[i]
32
        assert user_message.role == MessageRole.USER
33

34
        if i == 0:
35
            # make sure system prompt is included at the start
36
            str_message = f"{BOS} {B_INST} {system_message_str} "
37
        else:
38
            # end previous user-assistant interaction
39
            string_messages[-1] += f" {EOS}"
40
            # no need to include system prompt
41
            str_message = f"{BOS} {B_INST} "
42

43
        # include user message content
44
        str_message += f"{user_message.content} {E_INST}"
45

46
        if len(messages) > (i + 1):
47
            # if assistant message exists, add to str_message
48
            assistant_message = messages[i + 1]
49
            assert assistant_message.role == MessageRole.ASSISTANT
50
            str_message += f" {assistant_message.content}"
51

52
        string_messages.append(str_message)
53

54
    return "".join(string_messages)
55

56

57
def completion_to_prompt(completion: str, system_prompt: Optional[str] = None) -> str:
58
    system_prompt_str = system_prompt or DEFAULT_SYSTEM_PROMPT
59

60
    return (
61
        f"{BOS} {B_INST} {B_SYS} {system_prompt_str.strip()} {E_SYS} "
62
        f"{completion.strip()} {E_INST}"
63
    )
64