File size: 2,250 Bytes
d42b6d6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
import re
from pprint import pprint


class MessageComposer:
    def __init__(self):
        pass

    def merge(self, messages) -> str:
        self.messages = messages
        self.merged_str = ""
        for message in messages:
            role = message["role"]
            if role.lower() in ["system", "user"]:
                role_str = "me"
            elif role.lower() in ["assistant", "bot"]:
                role_str = "you"
            else:
                role_str = "unknown"
            content = message["content"]
            self.merged_str += f"`{role_str}`:\n{content}\n\n"
        self.merged_str += "`you`:\n"
        return self.merged_str

    def split(self, merged_str) -> list:
        self.messages = []
        self.merged_str = merged_str
        pattern = r"`(?P<role>me|you)`:\n(?P<content>.*)\n+"
        matches = re.finditer(pattern, self.merged_str, re.MULTILINE)
        for match in matches:
            role = match.group("role")

            if role == "me":
                role_str = "user"
            elif role == "you":
                role_str = "assistant"
            else:
                role_str = "unknown"

            self.messages.append(
                {
                    "role": role_str,
                    "content": match.group("content"),
                }
            )
        return self.messages


if __name__ == "__main__":
    composer = MessageComposer()
    messages = [
        {
            "role": "system",
            "content": "You are a LLM developed by OpenAI. Your name is GPT-4.",
        },
        {"role": "user", "content": "Hello, who are you?"},
        {"role": "assistant", "content": "I am a bot."},
        {"role": "user", "content": "What is your name?"},
        {"role": "assistant", "content": "My name is Bing."},
        {"role": "user", "content": "Tell me a joke."},
        {"role": "assistant", "content": "What is a robot's favorite type of music?"},
        {
            "role": "user",
            "content": "How many questions have I asked? Please list them.",
        },
    ]
    merged_str = composer.merge(messages)
    print(merged_str)
    pprint(composer.split(merged_str))
    print(composer.merge(composer.split(merged_str)))