Add predict update history.

Signed-off-by: ldwang <[email protected]>
FlagAI-Open · Oct 17, 2023 · 3875aa9 · 3875aa9
1 parent c8ef199
commit 3875aa9
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git a/flagai/model/aquila2/modeling_aquila.py b/flagai/model/aquila2/modeling_aquila.py
@@ -923,6 +923,7 @@ def predict(self, text, tokenizer=None,
                 sft=True, convo_template = "",
                 device = "cuda",
                 model_name="AquilaChat2-7B",
+                history=[],
                 **kwargs):
 
         vocab = tokenizer.get_vocab()
@@ -943,7 +944,7 @@ def predict(self, text, tokenizer=None,
             topk = 1
             temperature = 1.0
         if sft:
-            tokens = covert_prompt_to_input_ids_with_history(text, history=[], tokenizer=tokenizer, max_token=2048, convo_template=convo_template)
+            tokens = covert_prompt_to_input_ids_with_history(text, history=history, tokenizer=tokenizer, max_token=2048, convo_template=convo_template)
             tokens = torch.tensor(tokens)[None,].to(device)
         else :
             tokens = tokenizer.encode_plus(text)["input_ids"]
@@ -1031,6 +1032,11 @@ def predict(self, text, tokenizer=None,
 
             convert_tokens = convert_tokens[1:]
             probs = probs[1:]
+
+        # Update history
+        history.insert(0, ('USER', text))
+        history.insert(0, ('ASSISTANT', out))
+
         return out 
 
 @add_start_docstrings(

diff --git a/flagai/model/aquila2_hf/predict.py b/flagai/model/aquila2_hf/predict.py
@@ -333,6 +333,7 @@ def predict(model, text, tokenizer=None,
             sft=True, convo_template = "",
             device = "cuda",
             model_name="AquilaChat2-7B",
+            history=[],
             **kwargs):
 
     vocab = tokenizer.get_vocab()
@@ -352,7 +353,7 @@ def predict(model, text, tokenizer=None,
         topk = 1
         temperature = 1.0
     if sft:
-        tokens = covert_prompt_to_input_ids_with_history(text, history=[], tokenizer=tokenizer, max_token=2048, convo_template=convo_template)
+        tokens = covert_prompt_to_input_ids_with_history(text, history=history, tokenizer=tokenizer, max_token=2048, convo_template=convo_template)
         tokens = torch.tensor(tokens)[None,].to(device)
     else :
         tokens = tokenizer.encode_plus(text)["input_ids"]
@@ -433,4 +434,9 @@ def predict(model, text, tokenizer=None,
 
         convert_tokens = convert_tokens[1:]
         probs = probs[1:]
+
+    # Update history
+    history.insert(0, ('USER', text))
+    history.insert(0, ('ASSISTANT', out))
+
     return out