34 lines
1.0 KiB
Python
34 lines
1.0 KiB
Python
# The ChatML format
|
|
# Takes messages in and provides a text body to autocomplete on
|
|
|
|
def format(context, query, for_completion=False):
|
|
|
|
# ChatML format:
|
|
user_id = "user"
|
|
assistant_id = "assistant"
|
|
context_declaration = "<|im_start|>system\n"
|
|
message_declaration = "<|im_start|>{author}\n"
|
|
end_of_message = "<|im_end|>\n"
|
|
output = ""
|
|
|
|
if isinstance(query, str):
|
|
query = [{"author": "user", "body": query}]
|
|
|
|
if isinstance(query, list):
|
|
for message in query:
|
|
author = message["author"]
|
|
body = message["body"]
|
|
|
|
if "nickname" in message.keys():
|
|
nickname = message["nickname"]
|
|
author = nickname
|
|
|
|
output = f"{output}{message_declaration.format(author=author)}{body}{end_of_message}"
|
|
|
|
append = ""
|
|
if for_completion:
|
|
append = message_declaration.format(author=assistant_id)
|
|
|
|
output = f"""{context_declaration}{context}{end_of_message}{output}{append}"""
|
|
return output
|