[update]

main.py

@@ -36,7 +36,8 @@ def init_model(pretrained_model_name_or_path: str):
             pretrained_model_name_or_path,
             trust_remote_code=True,
             low_cpu_mem_usage=True,
-            torch_dtype=torch.bfloat16,
+            # torch_dtype=torch.bfloat16,
+            torch_dtype=torch.float16,
             device_map="auto",
             offload_folder="./offload",
             offload_state_dict=True,
@@ -45,7 +46,8 @@ def init_model(pretrained_model_name_or_path: str):
         if model.config.model_type == "chatglm":
             model = model.eval()
         else:
-            model = model.bfloat16().eval()
+            # model = model.bfloat16().eval()
+            model = model.eval()
 
         tokenizer = AutoTokenizer.from_pretrained(
             pretrained_model_name_or_path,
@@ -146,6 +148,8 @@ def chat_with_llm_streaming(question: str,
     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)
+        elif model.config.model_type == "llama2":
+            h_question = "Question: {}\n\nAnswer: ".format(h_question)
         utterances.append(h_question)
         utterances.append(h_answer)
     utterances.append(question)