add conv system in cli (#2098)

fastchat/conversation.py

@@ -257,7 +257,7 @@ def copy(self):
     def dict(self):
         return {
             "template_name": self.name,
-            "system": self.system_message,
+            "system_message": self.system_message,
             "roles": self.roles,
             "messages": self.messages,
             "offset": self.offset,

fastchat/serve/cli.py

@@ -208,6 +208,7 @@ def main(args):
             args.load_8bit,
             args.cpu_offloading,
             args.conv_template,
+            args.conv_system_msg,
             args.temperature,
             args.repetition_penalty,
             args.max_new_tokens,
@@ -238,6 +239,9 @@ def main(args):
     parser.add_argument(
         "--conv-template", type=str, default=None, help="Conversation prompt template."
     )
+    parser.add_argument(
+        "--conv-system-msg", type=str, default=None, help="Conversation system message."
+    )
     parser.add_argument("--temperature", type=float, default=0.7)
     parser.add_argument("--repetition_penalty", type=float, default=1.0)
     parser.add_argument("--max-new-tokens", type=int, default=512)

fastchat/serve/inference.py

@@ -1,7 +1,9 @@
 """Inference for FastChat models."""
 import abc
 import gc
+import json
 import math
+import os
 import sys
 import time
 from typing import Iterable, Optional, Dict
@@ -284,6 +286,7 @@ def chat_loop(
     load_8bit: bool,
     cpu_offloading: bool,
     conv_template: Optional[str],
+    conv_system_msg: Optional[str],
     temperature: float,
     repetition_penalty: float,
     max_new_tokens: int,
@@ -327,6 +330,8 @@ def new_chat():
             conv = get_conv_template(conv_template)
         else:
             conv = get_conversation_template(model_path)
+        if conv_system_msg is not None:
+            conv.set_system_message(conv_system_msg)
         return conv
 
     conv = None
@@ -343,12 +348,10 @@ def new_chat():
         if inp == "!!exit" or not inp:
             print("exit...")
             break
-
         elif inp == "!!reset":
             print("resetting...")
             conv = new_chat()
             continue
-
         elif inp.startswith("!!save"):
             args = inp.split(" ", 1)
 
@@ -362,14 +365,9 @@ def new_chat():
                 filename += ".json"
 
             print("saving...", filename)
-
-            import json
-
-            with open(filename, "w") as file:
-                json.dump(conv.dict(), file)
-
+            with open(filename, "w") as outfile:
+                json.dump(conv.dict(), outfile)
             continue
-
         elif inp.startswith("!!load"):
             args = inp.split(" ", 1)
 
@@ -379,8 +377,6 @@ def new_chat():
             else:
                 filename = args[1]
 
-            import os
-
             # Check if file exists and add .json if needed
             if not os.path.exists(filename):
                 if (not filename.endswith(".json")) and os.path.exists(
@@ -392,22 +388,15 @@ def new_chat():
                     continue
 
             print("loading...", filename)
-
-            import json
-
-            with open(filename, "r") as file:
-                new_conv = json.load(file)
+            with open(filename, "r") as infile:
+                new_conv = json.load(infile)
 
             conv = get_conv_template(new_conv["template_name"])
-            conv.system = new_conv["system"]
-            conv.roles = new_conv["roles"]
+            conv.set_system_message(new_conv["system_message"])
             conv.messages = new_conv["messages"]
-            conv.offset = new_conv["offset"]
-
             for message in conv.messages[conv.offset :]:
                 chatio.prompt_for_output(message[0])
                 chatio.print_output(message[1])
-
             continue
 
         conv.append_message(conv.roles[0], inp)