[update]add main

main.py

@@ -42,8 +42,11 @@ def init_model(pretrained_model_name_or_path: str):
             offload_state_dict=True,
             # load_in_4bit=True,
         )
-        model = model.to(device)
-        model = model.bfloat16().eval()
+        if model.config.model_type == "chatglm":
+            model = model.eval()
+        else:
+            model = model.to(device)
+            model = model.bfloat16().eval()
 
         tokenizer = AutoTokenizer.from_pretrained(
             pretrained_model_name_or_path,
@@ -79,18 +82,27 @@ def chat_with_llm_non_stream(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)
+    # input_ids
+    if model.config.model_type == "chatglm":
+        input_ids = []
+    else:
+        input_ids = [tokenizer.bos_token_id]
 
-    text_encoded = tokenizer.__call__(text_list, add_special_tokens=False)
-    batch_input_ids = text_encoded["input_ids"]
+    # history
+    utterances = list()
+    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)
+        utterances.append(h_question)
+        utterances.append(h_answer)
+    utterances.append(question)
+
+    encoded_utterances = tokenizer.__call__(utterances, add_special_tokens=False)
+    for encoded_utterance in encoded_utterances:
+        input_ids.extend(encoded_utterance)
+        if model.config.model_type == "chatglm":
+            input_ids.append(tokenizer.eos_token_id)
 
-    input_ids = [tokenizer.bos_token_id]
-    for input_ids_ in batch_input_ids:
-        input_ids.extend(input_ids_)
-        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)
 
@@ -122,31 +134,27 @@ def chat_with_llm_streaming(question: str,
 
     model, tokenizer = init_model(pretrained_model_name_or_path)
 
+    # input_ids
     if model.config.model_type == "chatglm":
         input_ids = []
     else:
         input_ids = [tokenizer.bos_token_id]
 
     # history
+    utterances = list()
     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)
+        utterances.append(h_question)
+        utterances.append(h_answer)
+    utterances.append(question)
 
-        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":
+    encoded_utterances = tokenizer.__call__(utterances, add_special_tokens=False)
+    for encoded_utterance in encoded_utterances:
+        input_ids.extend(encoded_utterance)
+        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)