Add aquila-v2 template.

Signed-off-by: ldwang <[email protected]>

flagai/model/aquila2/conversation.py

@@ -160,6 +160,8 @@ def get_conversation_template(model_path: str) -> Conversation:
     """Get the default conversation template."""
     if "aquila-v1" in model_path:
         return get_conv_template("aquila-v1")
+    elif "aquila-v2" in model_path:
+        return get_conv_template("aquila-v2")
     elif "aquila-chat" in model_path:
         return get_conv_template("aquila-chat")
     elif "aquila-legacy" in model_path:
@@ -227,6 +229,21 @@ def get_conversation_template(model_path: str) -> Conversation:
     )
 )
 
+register_conv_template(
+    Conversation(
+        name="aquila-v2",
+        system_message="A chat between a curious human and an artificial intelligence assistant. "
+        "The assistant gives helpful, detailed, and polite answers to the human's questions.\n\n",
+        roles=("<|startofpiece|>", "<|endofpiece|>", ""),
+        messages=(),
+        offset=0,
+        sep_style=SeparatorStyle.NO_COLON_TWO,
+        sep="",
+        sep2="</s>",
+        stop_str=["</s>", "<|endoftext|>"],
+    )
+)
+
 
 if __name__ == "__main__":
     print("aquila template:")
@@ -269,3 +286,13 @@ def get_conversation_template(model_path: str) -> Conversation:
 
     print("\n")
 
+    print("aquila-v2 template:")
+    conv = get_conv_template("aquila-v2")
+    conv.append_message(conv.roles[0], "Hello!")
+    conv.append_message(conv.roles[1], "Hi!")
+    conv.append_message(conv.roles[0], "How are you?")
+    conv.append_message(conv.roles[1], None)
+    print(conv.get_prompt())
+
+    print("\n")
+

flagai/model/aquila2_hf/predict.py

@@ -185,6 +185,8 @@ def get_conversation_template(model_path: str) -> Conversation:
     """Get the default conversation template."""
     if "aquila-v1" in model_path:
         return get_conv_template("aquila-v1")
+    elif "aquila-v2" in model_path:
+        return get_conv_template("aquila-v2")
     elif "aquila-chat" in model_path:
         return get_conv_template("aquila-chat")
     elif "aquila-legacy" in model_path:
@@ -252,6 +254,21 @@ def get_conversation_template(model_path: str) -> Conversation:
     )
 )
 
+register_conv_template(
+    Conversation(
+        name="aquila-v2",
+        system_message="A chat between a curious human and an artificial intelligence assistant. "
+        "The assistant gives helpful, detailed, and polite answers to the human's questions.\n\n",
+        roles=("<|startofpiece|>", "<|endofpiece|>", ""),
+        messages=(),
+        offset=0,
+        sep_style=SeparatorStyle.NO_COLON_TWO,
+        sep="",
+        sep2="</s>",
+        stop_str=["</s>", "<|endoftext|>"],
+    )
+)
+
 
 if __name__ == "__main__":
     print("aquila template:")
@@ -294,6 +311,17 @@ def get_conversation_template(model_path: str) -> Conversation:
 
     print("\n")
 
+    print("aquila-v2 template:")
+    conv = get_conv_template("aquila-v2")
+    conv.append_message(conv.roles[0], "Hello!")
+    conv.append_message(conv.roles[1], "Hi!")
+    conv.append_message(conv.roles[0], "How are you?")
+    conv.append_message(conv.roles[1], None)
+    print(conv.get_prompt())
+
+    print("\n")
+
+
 def set_random_seed(seed):
     """Set random seed for reproducability."""
     if seed is not None and seed > 0: