Skip to content

Commit

Permalink
feat: chatml-noforce-roles wrapper + cli fix (#738)
Browse files Browse the repository at this point in the history
* added new wrapper option to turn system style messages into system role messages

* added multirole wrapper

* unrelated issue with cli print due to previous PR (not showing core memory edits)
  • Loading branch information
cpacker authored Dec 29, 2023
1 parent b2d2357 commit cf2651d
Show file tree
Hide file tree
Showing 3 changed files with 13 additions and 2 deletions.
4 changes: 2 additions & 2 deletions memgpt/interface.py
Original file line number Diff line number Diff line change
Expand Up @@ -159,8 +159,8 @@ def printd_function_message(icon, msg, color=Fore.RED):
if match:
function_name = match.group(1)
function_args = match.group(2)
if function_name in ["archival_memory_insert", "archival_memory_search"]:
if function_name == "archival_memory_insert":
if function_name in ["archival_memory_insert", "archival_memory_search", "core_memory_replace", "core_memory_append"]:
if function_name in ["archival_memory_insert", "core_memory_append", "core_memory_replace"]:
print_function_message("🧠", f"updating memory with {function_name}")
elif function_name == "archival_memory_search":
print_function_message("🧠", f"searching memory with {function_name}")
Expand Down
9 changes: 9 additions & 0 deletions memgpt/local_llm/llm_chat_completion_wrappers/chatml.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@ def __init__(
assistant_prefix_extra='\n{\n "function":',
assistant_prefix_extra_first_message='\n{\n "function": "send_message",',
allow_custom_roles=True, # allow roles outside user/assistant
use_system_role_in_user=False, # use the system role on user messages that don't use "type: user_message"
# allow_function_role=True, # use function role for function replies?
allow_function_role=False, # use function role for function replies?
no_function_role_role="assistant", # if no function role, which role to use?
Expand All @@ -56,6 +57,7 @@ def __init__(

# role-based
self.allow_custom_roles = allow_custom_roles
self.use_system_role_in_user = use_system_role_in_user
self.allow_function_role = allow_function_role
# extras for when the function role is disallowed
self.no_function_role_role = no_function_role_role
Expand Down Expand Up @@ -198,6 +200,13 @@ def chat_completion_to_prompt(self, messages, functions, first_message=False):
role_str = message["name"].strip().lower() if (self.allow_custom_roles and "name" in message) else message["role"]
msg_str = self._compile_user_message(message)

if self.use_system_role_in_user:
try:
msg_json = json.loads(message["content"])
if msg_json["type"] != "user_message":
role_str = "system"
except:
pass
prompt += f"\n<|im_start|>{role_str}\n{msg_str.strip()}<|im_end|>"

elif message["role"] == "assistant":
Expand Down
2 changes: 2 additions & 0 deletions memgpt/local_llm/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,8 @@ def get_available_wrappers() -> dict:
# New chatml-based wrappers
"chatml": chatml.ChatMLInnerMonologueWrapper(),
"chatml-noforce": chatml.ChatMLOuterInnerMonologueWrapper(),
# "chatml-noforce-sysm": chatml.ChatMLOuterInnerMonologueWrapper(use_system_role_in_user=True),
"chatml-noforce-roles": chatml.ChatMLOuterInnerMonologueWrapper(use_system_role_in_user=True, allow_function_role=True),
# With extra hints
"chatml-hints": chatml.ChatMLInnerMonologueWrapper(assistant_prefix_hint=True),
"chatml-noforce-hints": chatml.ChatMLOuterInnerMonologueWrapper(assistant_prefix_hint=True),
Expand Down

0 comments on commit cf2651d

Please sign in to comment.