[update]add main

main.py

@@ -122,18 +122,31 @@ def chat_with_llm_streaming(question: str,
 
     model, tokenizer = init_model(pretrained_model_name_or_path)
 
-    text_list = list()
-    for pair in history:
-        text_list.extend(pair)
-    text_list.append(question)
-
-    text_encoded = tokenizer.__call__(text_list, add_special_tokens=False)
-    batch_input_ids = text_encoded["input_ids"]
-
-    input_ids = [tokenizer.bos_token_id]
-    for input_ids_ in batch_input_ids:
-        input_ids.extend(input_ids_)
+    if model.config.model_type == "chatglm":
+        input_ids = []
+    else:
+        input_ids = [tokenizer.bos_token_id]
+
+    # history
+    for idx, (h_question, h_answer) in enumerate(history):
+        if model.config.model_type == "chatglm":
+            h_question = "[Round {}]\n\n问：{}\n\n答：".format(idx, h_question)
+        h_question = tokenizer.__call__(h_question, add_special_tokens=False)
+        h_answer = tokenizer.__call__(h_answer, add_special_tokens=False)
+
+        input_ids.append(h_question)
+        if model.config.model_type != "chatglm":
+            input_ids.append(tokenizer.eos_token_id)
+        input_ids.append(h_answer)
+        if model.config.model_type != "chatglm":
+            input_ids.append(tokenizer.eos_token_id)
+
+    # question
+    question = tokenizer.__call__(question, add_special_tokens=False)
+    input_ids.append(question)
+    if model.config.model_type != "chatglm":
         input_ids.append(tokenizer.eos_token_id)
+
     input_ids = torch.tensor([input_ids], dtype=torch.long)
     input_ids = input_ids[:, -history_max_len:].to(device)