Using aquila-v2 template.

Signed-off-by: ldwang <[email protected]>
FlagAI-Open · Nov 28, 2023 · 5f35072 · 5f35072
1 parent 994ee63
commit 5f35072
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 2 deletions.
diff --git a/flagai/model/aquila2/conversation.py b/flagai/model/aquila2/conversation.py
@@ -240,7 +240,7 @@ def get_conversation_template(model_path: str) -> Conversation:
         sep_style=SeparatorStyle.NO_COLON_TWO,
         sep="",
         sep2="</s>",
-        stop_str=["</s>", "<|endoftext|>"],
+        stop_str=["</s>", "<|endoftext|>", "<|startofpiece|>", "<|endofpiece|>"],
     )
 )
 

diff --git a/flagai/model/aquila2/modeling_aquila.py b/flagai/model/aquila2/modeling_aquila.py
@@ -934,6 +934,7 @@ def predict(self, text, tokenizer=None,
 
         template_map = {"AquilaChat2-7B": "aquila-v1",
                         "AquilaChat2-34B": "aquila-legacy",
+                        "AquilaChat2-70B": "aquila-v2",
                         "AquilaChat2-7B-16K": "aquila",
                         "AquilaChat2-34B-16K": "aquila"}
         if not convo_template:

diff --git a/flagai/model/aquila2_hf/AquilaChat2-70B.py b/flagai/model/aquila2_hf/AquilaChat2-70B.py
@@ -0,0 +1,12 @@
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+model_info = "BAAI/AquilaChat2-70B"
+tokenizer = AutoTokenizer.from_pretrained(model_info, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(model_info, trust_remote_code=True, torch_dtype=torch.bfloat16)
+model.eval()
+text = "请给出10个要到北京旅游的理由。"
+from predict import predict
+out = predict(model, text, tokenizer=tokenizer, max_gen_len=200, top_p=0.95,
+              seed=1234, topk=100, temperature=0.9, sft=True,
+              model_name="AquilaChat2-70B")
+print(out)
diff --git a/flagai/model/aquila2_hf/predict.py b/flagai/model/aquila2_hf/predict.py
@@ -265,7 +265,7 @@ def get_conversation_template(model_path: str) -> Conversation:
         sep_style=SeparatorStyle.NO_COLON_TWO,
         sep="",
         sep2="</s>",
-        stop_str=["</s>", "<|endoftext|>"],
+        stop_str=["</s>", "<|endoftext|>", "<|startofpiece|>", "<|endofpiece|>"],
     )
 )
 
@@ -374,6 +374,7 @@ def predict(model, text, tokenizer=None,
 
     template_map = {"AquilaChat2-7B": "aquila-v1",
                     "AquilaChat2-34B": "aquila-legacy",
+                    "AquilaChat2-70B": "aquila-v2",
                     "AquilaChat2-7B-16K": "aquila",
                     "AquilaChat2-34B-16K": "aquila"}
     if not convo_template: