Update cyg_conversation.py

cyg_conversation.py

@@ -126,6 +126,30 @@ conv_templates = {
     "bair_v1": conv_bair_v1,
 }
 
+def covert_prompt_to_input_ids_with_history(text, history, tokenizer, max_token):
+    conv = default_conversation.copy()
+
+    conv.append_message(conv.roles[1], None)
+    conv.append_message(conv.roles[0], text)
+
+    example = tokenizer.encode_plus(f"{conv.get_prompt()}", None, max_length=None)['input_ids']
+
+    while(len(history) > 0 and (len(example) < max_token)):
+        tmp = history.pop()
+        if tmp[0] == 'ASSISTANT':
+            conv.append_message(conv.roles[1], tmp[1])
+        else:
+            conv.append_message(conv.roles[0], tmp[1])
+        example = tokenizer.encode_plus(f"{conv.get_prompt()}", None, max_length=None)['input_ids']
+
+    if len(example) >= max_token:
+        conv.messages.pop()
+    conv.messages = conv.messages[::-1]
+    print('model in:', conv.get_prompt())
+    example = tokenizer.encode_plus(f"{conv.get_prompt()}", None, max_length=None)['input_ids']
+    example = example[1:-1]
+
+    return example
 
 if __name__ == "__main__":
     print(default_conversation.get_prompt())